Merge master into main, resolved conflicts and updated LFS tracking

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/**/*.npy filter=lfs diff=lfs merge=lfs -text
+data/**/*.jsonl filter=lfs diff=lfs merge=lfs -text
+data/**/*.json filter=lfs diff=lfs merge=lfs -text
+assets/*.png filter=lfs diff=lfs merge=lfs -text
+data/**/*.npz filter=lfs diff=lfs merge=lfs -text

.huggingface.yaml

@@ -0,0 +1,3 @@
+# .huggingface.yaml
+sdk: streamlit  # or gradio
+app_file: ./app/ui.py

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Suman Adhya
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

app/ui_updated.py

@@ -0,0 +1,450 @@
+import streamlit as st
+import plotly.graph_objects as go
+import plotly.colors as pc
+import sys
+import os
+import base64
+import streamlit.components.v1 as components
+import html
+
+# Absolute path to the repo root (assuming `ui.py` is in /app)
+REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+sys.path.append(REPO_ROOT)
+ASSETS_DIR = os.path.join(REPO_ROOT, 'assets')
+DATA_DIR = os.path.join(REPO_ROOT, 'data')
+
+# Import functions from the backend
+from backend.inference.process_beta import (
+    load_beta_matrix,
+    get_top_words_over_time,
+    load_time_labels
+    )
+from backend.inference.word_selector import get_interesting_words, get_word_trend
+from backend.inference.indexing_utils import load_index
+from backend.inference.doc_retriever import (
+    load_length_stats,
+    get_yearly_counts_for_word,
+    deduplicate_docs,
+    get_all_documents_for_word_year,
+    highlight_words,
+    extract_snippet
+)
+from backend.llm_utils.summarizer import summarize_multiword_docs, ask_multiturn_followup
+from backend.llm_utils.label_generator import get_topic_labels
+from backend.llm.llm_router import get_llm, list_supported_models
+from backend.llm_utils.token_utils import estimate_k_max_from_word_stats
+
+def get_base64_image(image_path):
+    with open(image_path, "rb") as img_file:
+        return base64.b64encode(img_file.read()).decode()
+
+# --- Page Configuration ---
+st.set_page_config(
+    page_title="DTECT",
+    page_icon="🔍",
+    layout="wide"
+)
+
+# Sidebar branding and repo link
+st.sidebar.markdown(
+    """
+    <div style="text-align: center;">
+        <a href="https://github.com/dinb-ai/DTECT" target="_blank">
+            <img src="data:image/png;base64,{}" width="180" style="margin-bottom: 18px;">
+        </a>
+        <hr style="margin-bottom: 0;">
+    </div>
+    """.format(get_base64_image(os.path.join(ASSETS_DIR, 'Logo_light.png'))),
+    unsafe_allow_html=True
+)
+
+# 1. Sidebar: Model and Dataset Selection
+st.sidebar.title("Configuration")
+
+AVAILABLE_MODELS = ["DTM", "DETM", "CFDTM"]
+ENV_VAR_MAP = {
+    "OpenAI": "OPENAI_API_KEY",
+    "Anthropic": "ANTHROPIC_API_KEY",
+    "Gemini": "GEMINI_API_KEY",
+    "Mistral": "MISTRAL_API_KEY"
+}
+
+def list_datasets(data_dir):
+    return sorted([
+        name for name in os.listdir(data_dir)
+        if os.path.isdir(os.path.join(data_dir, name))
+    ])
+
+with st.sidebar.expander("Select Dataset & Topic Model", expanded=True):
+    datasets = list_datasets(DATA_DIR)
+    selected_dataset = st.selectbox("Dataset", datasets, help="Choose an available dataset.")
+    selected_model = st.selectbox("Model", AVAILABLE_MODELS, help="Select topic model architecture.")
+
+# Resolve paths
+dataset_path = os.path.join(DATA_DIR, selected_dataset)
+model_path = os.path.join(dataset_path, selected_model)
+docs_path = os.path.join(dataset_path, "docs.jsonl")
+vocab_path = os.path.join(dataset_path, "processed/vocab.txt")
+time2id_path = os.path.join(dataset_path, "processed/time2id.txt")
+index_path = os.path.join(dataset_path, "inverted_index.json")
+beta_path = os.path.join(model_path, "beta.npy")
+label_cache_path = os.path.join(model_path, "topic_label_cache.json")
+length_stats_path = os.path.join(dataset_path, "processed/length_stats.json")
+lemma_map_path = os.path.join(dataset_path, "processed/lemma_to_forms.json")
+
+with st.sidebar.expander("LLM Settings", expanded=True):
+    provider = st.selectbox("LLM Provider", options=list(ENV_VAR_MAP.keys()), help="Choose the LLM backend.")
+    available_models = list_supported_models(provider)
+    model = st.selectbox("LLM Model", options=available_models)
+    env_var = ENV_VAR_MAP[provider]
+    api_key = os.getenv(env_var)
+
+    if "llm_configured" not in st.session_state:
+        st.session_state.llm_configured = False
+
+    if api_key:
+        st.session_state.llm_configured = True
+    else:
+        st.session_state.llm_configured = False
+        with st.form(key="api_key_form"):
+            entered_key = st.text_input(f"Enter your {provider} API Key", type="password")
+            submitted = st.form_submit_button("Submit and Confirm")
+            if submitted:
+                if entered_key:
+                    os.environ[env_var] = entered_key
+                    api_key = entered_key
+                    st.session_state.llm_configured = True
+                    st.rerun()
+                else:
+                    st.warning("Please enter a key.")
+
+    if not st.session_state.llm_configured:
+        st.warning("Please configure your LLM settings in the sidebar.")
+        st.stop()
+
+    if api_key and not st.session_state.llm_configured:
+        st.session_state.llm_configured = True
+
+    if not api_key:
+        st.session_state.llm_configured = False
+
+    if not st.session_state.llm_configured:
+        st.warning("Please configure your LLM settings in the sidebar.")
+        st.stop()
+
+# Initialize LLM with the provided key
+llm = get_llm(provider=provider, model=model, api_key=api_key)
+
+# 3. Load Data
+@st.cache_resource
+def load_resources(beta_path, vocab_path, docs_path, index_path, time2id_path, length_stats_path, lemma_map_path):
+    beta, vocab = load_beta_matrix(beta_path, vocab_path)
+    index, docs, lemma_to_forms = load_index(docs_file_path=docs_path, vocab=vocab, index_path=index_path, lemma_map_path=lemma_map_path)
+    time_labels = load_time_labels(time2id_path)
+    length_stats = load_length_stats(length_stats_path)
+    return beta, vocab, index, docs, lemma_to_forms, time_labels, length_stats
+
+# --- Main Title and Paper-aligned Intro ---
+st.markdown("""# 🔍 DTECT: Dynamic Topic Explorer & Context Tracker""")
+
+# --- Load resources ---
+try:
+    beta, vocab, index, docs, lemma_to_forms, time_labels, length_stats = load_resources(
+        beta_path,
+        vocab_path,
+        docs_path,
+        index_path,
+        time2id_path,
+        length_stats_path,
+        lemma_map_path
+    )
+except FileNotFoundError as e:
+    st.error(f"Missing required file: {e}")
+    st.stop()
+except Exception as e:
+    st.error(f"Failed to load data: {str(e)}")
+    st.stop()
+
+timestamps = list(range(len(time_labels)))
+num_topics = beta.shape[1]
+# Estimate max_k based on document length stats and selected LLM
+suggested_max_k = estimate_k_max_from_word_stats(length_stats.get("avg_len"), model_name=model, provider=provider)
+
+
+# ==============================================================================
+# 1. 🏷 TOPIC LABELING
+# ==============================================================================
+st.markdown("## 1️⃣ 🏷️ Topic Labeling")
+st.info("Topics are automatically labeled using LLMs by analyzing their temporal word distributions.")
+
+topic_labels = get_topic_labels(beta, vocab, time_labels, llm, label_cache_path)
+topic_options = list(topic_labels.values())
+selected_topic_label = st.selectbox("Select a Topic", topic_options, help="LLM-generated topic label")
+label_to_topic = {v: k for k, v in topic_labels.items()}
+selected_topic = label_to_topic[selected_topic_label]
+
+# ==============================================================================
+# 2. 💡 INFORMATIVE WORD DETECTION & 📊 TREND VISUALIZATION
+# ==============================================================================
+st.markdown("---")
+st.markdown("## 2️⃣ 💡 Informative Word Detection & 📊 Trend Visualization")
+st.info("Explore top/interesting words for each topic, and visualize their trends over time.")
+
+top_n_words = st.slider("Number of Top Words per Topic", min_value=5, max_value=500, value=10)
+top_words = get_top_words_over_time(
+    beta=beta,
+    vocab=vocab,
+    topic_id=selected_topic,
+    top_n=top_n_words
+)
+
+st.write(f"### Top {top_n_words} Words for Topic '{selected_topic_label}' (Ranked):")
+scrollable_top_words = "<div style='max-height: 200px; overflow-y: auto; padding: 0 10px;'>"
+words_per_col = (top_n_words + 3) // 4
+columns = [top_words[i:i+words_per_col] for i in range(0, len(top_words), words_per_col)]
+scrollable_top_words += "<div style='display: flex; gap: 20px;'>"
+word_rank = 1
+for col in columns:
+    scrollable_top_words += "<div style='flex: 1;'>"
+    for word in col:
+        scrollable_top_words += f"<div style='margin-bottom: 4px;'>{word_rank}. {word}</div>"
+        word_rank += 1
+    scrollable_top_words += "</div>"
+scrollable_top_words += "</div></div>"
+st.markdown(scrollable_top_words, unsafe_allow_html=True)
+
+st.markdown("<div style='margin-top: 18px;'></div>", unsafe_allow_html=True)
+
+if st.button("💡 Suggest Informative Words", key="suggest_topic_words"):
+    top_words = get_top_words_over_time(
+        beta=beta,
+        vocab=vocab,
+        topic_id=selected_topic,
+        top_n=top_n_words
+    )
+    interesting_words = get_interesting_words(beta, vocab, topic_id=selected_topic, restrict_to=top_words)
+    st.session_state.interesting_words = interesting_words
+    st.session_state.selected_words = interesting_words[:15]  # pre-fill multiselect
+    styled_words = " ".join([
+        f"<span style='background-color:#e0f7fa; color:#004d40; font-weight:500; padding:4px 8px; margin:4px; border-radius:8px; display:inline-block;'>{w}</span>"
+        for w in interesting_words
+    ])
+    st.markdown(
+        f"**Top Informative Words from Topic '{selected_topic_label}':**<br>{styled_words}",
+        unsafe_allow_html=True
+    )
+
+st.markdown("#### 📈 Plot Word Trends Over Time")
+all_word_options = vocab
+interesting_words = st.session_state.get("interesting_words", [])
+
+if "selected_words" not in st.session_state:
+    st.session_state.selected_words = interesting_words[:15]  # initial default
+
+selected_words = st.multiselect(
+    "Select words to visualize trends",
+    options=all_word_options,
+    default=st.session_state.selected_words,
+    key="selected_words"
+)
+if selected_words:
+    fig = go.Figure()
+    color_cycle = pc.qualitative.Plotly
+    for i, word in enumerate(selected_words):
+        trend = get_word_trend(beta, vocab, word, topic_id=selected_topic)
+        color = color_cycle[i % len(color_cycle)]
+        fig.add_trace(go.Scatter(
+            x=time_labels,
+            y=trend,
+            name=word,
+            line=dict(color=color),
+            legendgroup=word,
+            showlegend=True
+        ))
+    fig.update_layout(title="", xaxis_title="Year", yaxis_title="Importance")
+    st.plotly_chart(fig, use_container_width=True)
+
+# ==============================================================================
+# 3. 🔍 DOCUMENT RETRIEVAL & 📃 SUMMARIZATION
+# ==============================================================================
+st.markdown("---")
+st.markdown("## 3️⃣ 🔍 Document Retrieval & 📃 Summarization")
+st.info("Retrieve and summarize documents matching selected words and years.")
+
+if selected_words:
+    st.markdown("#### 📊 Document Frequency Over Time")
+    selected_words_for_counts = st.multiselect(
+        "Select word(s) to show document frequencies over time",
+        options=selected_words,
+        default=selected_words[:3],
+        key="word_counts_multiselect"
+    )
+
+    if selected_words_for_counts:
+        color_cycle = pc.qualitative.Set2
+        bar_fig = go.Figure()
+        for i, word in enumerate(selected_words_for_counts):
+            doc_years, doc_counts = get_yearly_counts_for_word(index=index, word=word)
+            bar_fig.add_trace(go.Bar(
+                x=doc_years,
+                y=doc_counts,
+                name=word,
+                marker_color=color_cycle[i % len(color_cycle)],
+                opacity=0.85
+            ))
+        bar_fig.update_layout(
+            barmode="group",
+            title="Document Frequency Over Time",
+            xaxis_title="Year",
+            yaxis_title="Document Count",
+            xaxis=dict(
+                tickmode='linear',
+                dtick=1,
+                tickformat='d'
+            ),
+            bargap=0.2
+        )
+        st.plotly_chart(bar_fig, use_container_width=True)
+
+    st.markdown("#### 📄 Inspect Documents for Word-Year Pairs")
+    # selected_year = st.slider("Select year", min_value=int(time_labels[0]), max_value=int(time_labels[-1]), key="inspect_year_slider")
+    selected_year = st.selectbox(
+        "Select year",
+        options=time_labels, # Use the list of available time labels (years)
+        index=0, # Default to the first year in the list
+        key="inspect_year_selectbox"
+    )
+    collected_docs_raw = []
+    for word in selected_words_for_counts:
+        docs_for_word_year = get_all_documents_for_word_year(
+            index=index,
+            docs_file_path=docs_path,
+            word=word,
+            year=selected_year
+        )
+        for doc in docs_for_word_year:
+            doc["__word__"] = word
+        collected_docs_raw.extend(docs_for_word_year)
+
+    if collected_docs_raw:
+        st.session_state.collected_deduplicated_docs = deduplicate_docs(collected_docs_raw)
+        st.write(f"Found {len(collected_docs_raw)} matching documents, {len(st.session_state.collected_deduplicated_docs)} after deduplication.")
+
+        html_blocks = ""
+        for doc in st.session_state.collected_deduplicated_docs:
+            word = doc["__word__"]
+            full_text = html.escape(doc["text"])
+            snippet_text = extract_snippet(doc["text"], word)
+            highlighted_snippet = highlight_words(
+                snippet_text,
+                query_words=selected_words_for_counts,
+                lemma_to_forms=lemma_to_forms
+            )
+            html_blocks += f"""
+            <div style="margin-bottom: 14px; padding: 10px; background-color: #fffbe6; border: 1px solid #f0e6cc; border-radius: 6px;">
+            <div style="color: #333;"><strong>Match:</strong> {word} | <strong>Doc ID:</strong> {doc['id']} | <strong>Timestamp:</strong> {doc['timestamp']}</div>
+            <div style="margin-top: 4px; color: #444;"><em>Snippet:</em> {highlighted_snippet}</div>
+            <details style="margin-top: 4px;">
+                <summary style="cursor: pointer; color: #007acc;">Show full document</summary>
+                <pre style="white-space: pre-wrap; color: #111; background-color: #fffef5; padding: 8px; border: 1px solid #f0e6cc; border-radius: 4px;">{full_text}</pre>
+            </details>
+            </div>
+            """
+        min_height = 120
+        max_height = 700
+        per_doc_height = 130
+        dynamic_height = min_height + per_doc_height * max(len(st.session_state.collected_deduplicated_docs) - 1, 0)
+        container_height = min(dynamic_height, max_height)
+        scrollable_html = f"""
+            <div style="overflow-y: auto; padding: 10px; 
+                        border: 1px solid #f0e6cc; border-radius: 6px; 
+                        background-color: #fffbe6; color: #222;
+                        margin-bottom: 0;">
+                {html_blocks}
+            </div>
+        """
+        components.html(scrollable_html, height=container_height, scrolling=True)
+    else:
+        st.warning("No documents found for the selected words and year.")
+
+# ==============================================================================
+# 4. 💬 CHAT ASSISTANT (Summary & Follow-up)
+# ==============================================================================
+st.markdown("---")
+st.markdown("## 4️⃣ 💬 Chat Assistant")
+st.info("Generate summaries from the inspected documents and ask follow-up questions.")
+
+if "summary" not in st.session_state:
+    st.session_state.summary = None
+if "context_for_followup" not in st.session_state:
+    st.session_state.context_for_followup = ""
+if "followup_history" not in st.session_state:
+    st.session_state.followup_history = []
+
+# MMR K selection
+st.markdown(f"**Max documents for summarization (k):**")
+st.markdown(f"The suggested maximum number of documents for summarization (k) based on the average document length and the selected LLM is **{suggested_max_k}**.")
+mmr_k = st.slider(
+    "Select the maximum number of documents (k) for MMR (Maximum Marginal Relevance) selection for summarization.",
+    min_value=1,
+    max_value=20, # Set a reasonable max for k, can be adjusted
+    value=min(suggested_max_k, 20), # Use suggested_max_k as default, capped at 20
+    help="This value determines how many relevant and diverse documents will be selected for summarization."
+)
+
+if st.button("📃 Summarize These Documents"):
+    if st.session_state.get("collected_deduplicated_docs"):
+        st.session_state.summary = None
+        st.session_state.context_for_followup = ""
+        st.session_state.followup_history = []
+        with st.spinner("Selecting and summarizing documents..."):
+            summary, mmr_docs = summarize_multiword_docs(
+                selected_words_for_counts,
+                selected_year,
+                st.session_state.collected_deduplicated_docs,
+                llm,
+                k=mmr_k
+            )
+            st.session_state.summary = summary
+            st.session_state.context_for_followup = "\n".join(
+                f"Document {i+1}:\n{doc.page_content.strip()}" for i, doc in enumerate(mmr_docs)
+            )
+            st.session_state.followup_history.append(
+                {"role": "user", "content": f"Please summarize the context of the words '{', '.join(selected_words_for_counts)}' in {selected_year} based on the provided documents."}
+            )
+            st.session_state.followup_history.append(
+                {"role": "assistant", "content": st.session_state.summary}
+            )
+        st.success(f"✅ Summary generated from {len(mmr_docs)} MMR-selected documents.")
+    else:
+        st.warning("⚠️ No documents collected to summarize. Please inspect some documents first.")
+
+if st.session_state.summary:
+    st.markdown(f"**Summary for words `{', '.join(selected_words_for_counts)}` in `{selected_year}`:**")
+    st.write(st.session_state.summary)
+
+    if st.checkbox("💬 Ask follow-up questions about this summary", key="enable_followup"):
+        with st.expander("View the documents used for this conversation"):
+            st.text_area("Context Documents", st.session_state.context_for_followup, height=200)
+        st.info("Ask a question based on the summary and the documents above.")
+        for msg in st.session_state.followup_history[2:]:
+            with st.chat_message(msg["role"], avatar="🧑" if msg["role"] == "user" else "🤖"):
+                st.markdown(msg["content"])
+        if user_query := st.chat_input("Ask a follow-up question..."):
+            with st.chat_message("user", avatar="🧑"):
+                st.markdown(user_query)
+            st.session_state.followup_history.append({"role": "user", "content": user_query})
+            with st.spinner("Thinking..."):
+                followup_response = ask_multiturn_followup(
+                    history=st.session_state.followup_history,
+                    question=user_query,
+                    llm=llm,
+                    context_texts=st.session_state.context_for_followup
+                )
+            st.session_state.followup_history.append({"role": "assistant", "content": followup_response})
+            if followup_response.startswith("[Error"):
+                st.error(followup_response)
+            else:
+                with st.chat_message("assistant", avatar="🤖"):
+                    st.markdown(followup_response)
+            st.rerun()

backend/__init__.py

@@ -0,0 +1,81 @@
+# === Inference components ===
+from .inference.process_beta import (
+    load_beta_matrix,
+    get_top_words_at_time,
+    get_top_words_over_time,
+    load_time_labels
+    )
+
+from .inference.indexing_utils import load_index
+from .inference.word_selector import (
+    get_interesting_words,
+    get_word_trend
+)
+from .inference.peak_detector import detect_peaks
+from .inference.doc_retriever import (
+    load_length_stats,
+    get_yearly_counts_for_word,
+    get_all_documents_for_word_year,
+    deduplicate_docs,
+    extract_snippet,
+    highlight,
+    get_docs_by_ids,
+)
+
+# === LLM components ===
+from .llm_utils.label_generator import label_topic_temporal, get_topic_labels
+from .llm_utils.token_utils import (
+    get_token_limit_for_model,
+    count_tokens,
+    estimate_avg_tokens_per_doc,
+    estimate_max_k,
+    estimate_max_k_fast
+    )
+from .llm_utils.summarizer import (
+    summarize_docs,
+    summarize_multiword_docs,
+    ask_multiturn_followup
+)
+from .llm.llm_router import (
+    list_supported_models,
+    get_llm
+)
+
+# === Dataset utilities ===
+from .datasets import dynamic_dataset
+from .datasets import preprocess
+from .datasets.utils import logger, _utils
+from .datasets.data import file_utils, download
+
+# === Evaluation ===
+from .evaluation.CoherenceModel_ttc import CoherenceModel_ttc
+from .evaluation.eval import TopicQualityAssessor
+
+# === Models ===
+from .models.DETM import DETM
+from .models.DTM_trainer import DTMTrainer
+from .models.CFDTM.CFDTM import CFDTM
+from .models.dynamic_trainer import DynamicTrainer
+
+__all__ = [
+    # Inference
+    "load_beta_matrix", "load_time_labels", "get_top_words_at_time", "get_top_words_over_time",
+    "load_index", "get_interesting_words", "get_word_trend", "detect_peaks",
+    "load_length_stats", "get_yearly_counts_for_word", "get_all_documents_for_word_year",
+    "deduplicate_docs", "extract_snippet", "highlight", "get_docs_by_ids",
+
+    # LLM
+    "summarize_docs", "summarize_multiword_docs", "ask_multiturn_followup",
+    "get_token_limit_for_model", "list_supported_models", "get_llm",
+    "label_topic_temporal", "get_topic_labels", "count_tokens",
+    "estimate_avg_tokens_per_doc", "estimate_max_k", "estimate_max_k_fast",
+
+    # Dataset
+    "dynamic_dataset", "preprocess", "logger","_utils", "file_utils", "download", 
+
+    # Evaluation
+    "CoherenceModel_ttc", "TopicQualityAssessor",
+
+    # Models
+    "DETM", "DTMTrainer", "CFDTM", "DynamicTrainer"
+]

backend/datasets/_preprocess.py

@@ -0,0 +1,447 @@
+import os
+import re
+import string
+import gensim.downloader
+from collections import Counter
+import numpy as np
+import scipy.sparse
+from tqdm import tqdm
+from sklearn.feature_extraction.text import CountVectorizer
+
+from backend.datasets.data import file_utils
+from backend.datasets.utils._utils import get_stopwords_set
+from backend.datasets.utils.logger import Logger
+import json
+import nltk
+from nltk.stem import WordNetLemmatizer
+
+logger = Logger("WARNING")
+
+try:
+    nltk.data.find('corpora/wordnet')
+except LookupError:
+    nltk.download('wordnet', quiet=True)
+try:
+    nltk.data.find('corpora/omw-1.4')
+except LookupError:
+    nltk.download('omw-1.4', quiet=True)
+
+# compile some regexes
+punct_chars = list(set(string.punctuation) - set("'"))
+punct_chars.sort()
+punctuation = ''.join(punct_chars)
+replace = re.compile('[%s]' % re.escape(punctuation))
+alpha = re.compile('^[a-zA-Z_]+$')
+alpha_or_num = re.compile('^[a-zA-Z_]+|[0-9_]+$')
+alphanum = re.compile('^[a-zA-Z0-9_]+$')
+
+
+class Tokenizer:
+    def __init__(self,
+                 stopwords="English",
+                 keep_num=False,
+                 keep_alphanum=False,
+                 strip_html=False,
+                 no_lower=False,
+                 min_length=3,
+                 lemmatize=True,
+                ):
+        self.keep_num = keep_num
+        self.keep_alphanum = keep_alphanum
+        self.strip_html = strip_html
+        self.lower = not no_lower
+        self.min_length = min_length
+
+        self.stopword_set = get_stopwords_set(stopwords)
+        
+        self.lemmatize = lemmatize
+        if lemmatize:
+            self.lemmatizer = WordNetLemmatizer()
+
+    def clean_text(self, text, strip_html=False, lower=True, keep_emails=False, keep_at_mentions=False):
+        # remove html tags
+        if strip_html:
+            text = re.sub(r'<[^>]+>', '', text)
+        else:
+            # replace angle brackets
+            text = re.sub(r'<', '(', text)
+            text = re.sub(r'>', ')', text)
+        # lower case
+        if lower:
+            text = text.lower()
+        # eliminate email addresses
+        if not keep_emails:
+            text = re.sub(r'\S+@\S+', ' ', text)
+        # eliminate @mentions
+        if not keep_at_mentions:
+            text = re.sub(r'\s@\S+', ' ', text)
+        # replace underscores with spaces
+        text = re.sub(r'_', ' ', text)
+        # break off single quotes at the ends of words
+        text = re.sub(r'\s\'', ' ', text)
+        text = re.sub(r'\'\s', ' ', text)
+        # remove periods
+        text = re.sub(r'\.', '', text)
+        # replace all other punctuation (except single quotes) with spaces
+        text = replace.sub(' ', text)
+        # remove single quotes
+        text = re.sub(r'\'', '', text)
+        # replace all whitespace with a single space
+        text = re.sub(r'\s', ' ', text)
+        # strip off spaces on either end
+        text = text.strip()
+        return text
+
+    def tokenize(self, text):
+        text = self.clean_text(text, self.strip_html, self.lower)
+        tokens = text.split()
+
+        tokens = ['_' if t in self.stopword_set else t for t in tokens]
+
+        # remove tokens that contain numbers
+        if not self.keep_alphanum and not self.keep_num:
+            tokens = [t if alpha.match(t) else '_' for t in tokens]
+
+        # or just remove tokens that contain a combination of letters and numbers
+        elif not self.keep_alphanum:
+            tokens = [t if alpha_or_num.match(t) else '_' for t in tokens]
+
+        # drop short tokens
+        if self.min_length > 0:
+            tokens = [t if len(t) >= self.min_length else '_' for t in tokens]
+            
+        if getattr(self, "lemmatize", False):
+            tokens = [self.lemmatizer.lemmatize(t) if t != '_' else t for t in tokens]
+
+        unigrams = [t for t in tokens if t != '_']
+        return unigrams
+
+
+def make_word_embeddings(vocab):
+    glove_vectors = gensim.downloader.load('glove-wiki-gigaword-200')
+    word_embeddings = np.zeros((len(vocab), glove_vectors.vectors.shape[1]))
+
+    num_found = 0
+
+    try:
+        key_word_list = glove_vectors.index_to_key
+    except:
+        key_word_list = glove_vectors.index2word
+
+    for i, word in enumerate(tqdm(vocab, desc="loading word embeddings")):
+        if word in key_word_list:
+            word_embeddings[i] = glove_vectors[word]
+            num_found += 1
+
+    logger.info(f'number of found embeddings: {num_found}/{len(vocab)}')
+
+    return scipy.sparse.csr_matrix(word_embeddings)
+
+
+class Preprocess:
+    def __init__(self,
+                 tokenizer=None,
+                 test_sample_size=None,
+                 test_p=0.2,
+                 stopwords="English",
+                 min_doc_count=0,
+                 max_doc_freq=1.0,
+                 keep_num=False,
+                 keep_alphanum=False,
+                 strip_html=False,
+                 no_lower=False,
+                 min_length=3,
+                 min_term=0,
+                 vocab_size=None,
+                 seed=42,
+                 verbose=True,
+                 lemmatize=True,
+                ):
+        """
+        Args:
+            test_sample_size:
+                Size of the test set.
+            test_p:
+                Proportion of the test set. This helps sample the train set based on the size of the test set.
+            stopwords:
+                List of stopwords to exclude.
+            min-doc-count:
+                Exclude words that occur in less than this number of documents.
+            max_doc_freq:
+                Exclude words that occur in more than this proportion of documents.
+            keep-num:
+                Keep tokens made of only numbers.
+            keep-alphanum:
+                Keep tokens made of a mixture of letters and numbers.
+            strip_html:
+                Strip HTML tags.
+            no-lower:
+                Do not lowercase text
+            min_length:
+                Minimum token length.
+            min_term:
+                Minimum term number
+            vocab-size:
+                Size of the vocabulary (by most common in the union of train and test sets, following above exclusions)
+            seed:
+                Random integer seed (only relevant for choosing test set)
+            lemmatize:
+                Whether to apply lemmatization to the tokens.
+        """
+
+        self.test_sample_size = test_sample_size
+        self.min_doc_count = min_doc_count
+        self.max_doc_freq = max_doc_freq
+        self.min_term = min_term
+        self.test_p = test_p
+        self.vocab_size = vocab_size
+        self.seed = seed
+
+        if tokenizer is not None:
+            self.tokenizer = tokenizer
+        else:
+            self.tokenizer = Tokenizer(
+                stopwords,
+                keep_num,
+                keep_alphanum,
+                strip_html,
+                no_lower,
+                min_length,
+                lemmatize=lemmatize
+            ).tokenize
+
+        if verbose:
+            logger.set_level("DEBUG")
+        else:
+            logger.set_level("WARNING")
+
+    def parse(self, texts, vocab):
+        if not isinstance(texts, list):
+            texts = [texts]
+
+        vocab_set = set(vocab)
+        parsed_texts = list()
+        for i, text in enumerate(tqdm(texts, desc="parsing texts")):
+            tokens = self.tokenizer(text)
+            tokens = [t for t in tokens if t in vocab_set]
+            parsed_texts.append(" ".join(tokens))
+
+        vectorizer = CountVectorizer(vocabulary=vocab, tokenizer=lambda x: x.split())
+        sparse_bow = vectorizer.fit_transform(parsed_texts)
+        return parsed_texts, sparse_bow
+
+    def preprocess_jsonlist(self, dataset_dir, label_name=None, use_partition=True):
+        if use_partition:
+            train_items = file_utils.read_jsonlist(os.path.join(dataset_dir, 'train.jsonlist'))
+            test_items = file_utils.read_jsonlist(os.path.join(dataset_dir, 'test.jsonlist'))
+        else:
+            raw_path = os.path.join(dataset_dir, 'docs.jsonl')
+            with open(raw_path, 'r', encoding='utf-8') as f:
+                train_items = [json.loads(line.strip()) for line in f if line.strip()]
+            test_items = []
+
+        logger.info(f"Found training documents {len(train_items)} testing documents {len(test_items)}")
+
+        # Initialize containers
+        raw_train_texts, train_labels, raw_train_times = [], [], []
+        raw_test_texts, test_labels, raw_test_times = [], [], []
+
+        # Process train items
+        for item in train_items:
+            raw_train_texts.append(item['text'])
+            raw_train_times.append(str(item['timestamp']))
+            if label_name and label_name in item:
+                train_labels.append(item[label_name])
+
+        # Process test items
+        for item in test_items:
+            raw_test_texts.append(item['text'])
+            raw_test_times.append(str(item['timestamp']))
+            if label_name and label_name in item:
+                test_labels.append(item[label_name])
+
+        # Create and apply time2id mapping
+        all_times = sorted(set(raw_train_times + raw_test_times))
+        time2id = {t: i for i, t in enumerate(all_times)}
+
+        train_times = np.array([time2id[t] for t in raw_train_times], dtype=np.int32)
+        test_times = np.array([time2id[t] for t in raw_test_times], dtype=np.int32) if raw_test_times else None
+
+        # Preprocess and get sample indices
+        rst = self.preprocess(raw_train_texts, train_labels, raw_test_texts, test_labels)
+        train_idx = rst.get("train_idx")
+        test_idx = rst.get("test_idx")
+
+        # Add filtered timestamps to result for saving later
+        rst["train_times"] = train_times[train_idx]
+        if test_times is not None and test_idx is not None:
+            rst["test_times"] = test_times[test_idx]
+
+        # Add time2id to result dict
+        rst["time2id"] = time2id
+
+        return rst
+
+
+    def convert_labels(self, train_labels, test_labels):
+        if train_labels:
+            label_list = list(set(train_labels).union(set(test_labels)))
+            label_list.sort()
+            n_labels = len(label_list)
+            label2id = dict(zip(label_list, range(n_labels)))
+
+            logger.info(f"label2id: {label2id}")
+
+            train_labels = [label2id[label] for label in train_labels]
+
+            if test_labels:
+                test_labels = [label2id[label] for label in test_labels]
+
+        return train_labels, test_labels
+
+    def preprocess(
+            self,
+            raw_train_texts,
+            train_labels=None,
+            raw_test_texts=None,
+            test_labels=None,
+            pretrained_WE=True
+        ):
+        np.random.seed(self.seed)
+
+        train_texts = list()
+        test_texts = list()
+        word_counts = Counter()
+        doc_counts_counter = Counter()
+
+        train_labels, test_labels = self.convert_labels(train_labels, test_labels)
+
+        for text in tqdm(raw_train_texts, desc="loading train texts"):
+            tokens = self.tokenizer(text)
+            word_counts.update(tokens)
+            doc_counts_counter.update(set(tokens))
+            parsed_text = ' '.join(tokens)
+            train_texts.append(parsed_text)
+
+        if raw_test_texts:
+            for text in tqdm(raw_test_texts, desc="loading test texts"):
+                tokens = self.tokenizer(text)
+                word_counts.update(tokens)
+                doc_counts_counter.update(set(tokens))
+                parsed_text = ' '.join(tokens)
+                test_texts.append(parsed_text)
+
+        words, doc_counts = zip(*doc_counts_counter.most_common())
+        doc_freqs = np.array(doc_counts) / float(len(train_texts) + len(test_texts))
+
+        vocab = [word for i, word in enumerate(words) if doc_counts[i] >= self.min_doc_count and doc_freqs[i] <= self.max_doc_freq]
+
+        # filter vocabulary
+        if self.vocab_size is not None:
+            vocab = vocab[:self.vocab_size]
+
+        vocab.sort()
+
+        train_idx = [i for i, text in enumerate(train_texts) if len(text.split()) >= self.min_term]
+        train_idx = np.asarray(train_idx)
+
+        if raw_test_texts is not None:
+            test_idx = [i for i, text in enumerate(test_texts) if len(text.split()) >= self.min_term]
+            test_idx = np.asarray(test_idx)
+        else:
+            test_idx = None
+
+        # randomly sample
+        if self.test_sample_size and raw_test_texts is not None:
+            logger.info("sample train and test sets...")
+
+            train_num = len(train_idx)
+            test_num = len(test_idx)
+            test_sample_size = min(test_num, self.test_sample_size)
+            train_sample_size = int((test_sample_size / self.test_p) * (1 - self.test_p))
+            if train_sample_size > train_num:
+                test_sample_size = int((train_num / (1 - self.test_p)) * self.test_p)
+                train_sample_size = train_num
+
+            train_idx = train_idx[np.sort(np.random.choice(train_num, train_sample_size, replace=False))]
+            test_idx = test_idx[np.sort(np.random.choice(test_num, test_sample_size, replace=False))]
+
+            logger.info(f"sampled train size: {len(train_idx)}")
+            logger.info(f"sampled test size: {len(test_idx)}")
+
+        train_texts, train_bow = self.parse([train_texts[i] for i in train_idx], vocab)
+
+        rst = {
+            'vocab': vocab,
+            'train_bow': train_bow,
+            "train_texts": train_texts,
+            "train_idx": train_idx,  # <--- NEW: indices of kept train samples
+        }
+
+        if train_labels:
+            rst['train_labels'] = np.asarray(train_labels)[train_idx]
+
+        logger.info(f"Real vocab size: {len(vocab)}")
+        logger.info(f"Real training size: {len(train_texts)} \t avg length: {rst['train_bow'].sum() / len(train_texts):.3f}")
+
+        if raw_test_texts:
+            rst['test_texts'], rst['test_bow'] = self.parse(np.asarray(test_texts)[test_idx].tolist(), vocab)
+            rst["test_idx"] = test_idx  # <--- NEW: indices of kept test samples
+
+            if test_labels:
+                rst['test_labels'] = np.asarray(test_labels)[test_idx]
+
+            logger.info(f"Real testing size: {len(rst['test_texts'])} \t avg length: {rst['test_bow'].sum() / len(rst['test_texts']):.3f}")
+
+        if pretrained_WE:
+            rst['word_embeddings'] = make_word_embeddings(vocab)
+
+        return rst
+
+    def save(
+        self,
+        output_dir,
+        vocab,
+        train_texts,
+        train_bow,
+        word_embeddings=None,
+        train_labels=None,
+        test_texts=None,
+        test_bow=None,
+        test_labels=None,
+        train_times=None,
+        test_times=None,
+        time2id=None  # <-- new parameter
+    ):
+        file_utils.make_dir(output_dir)
+
+        file_utils.save_text(vocab, f"{output_dir}/vocab.txt")
+        file_utils.save_text(train_texts, f"{output_dir}/train_texts.txt")
+        scipy.sparse.save_npz(f"{output_dir}/train_bow.npz", scipy.sparse.csr_matrix(train_bow))
+        
+        if word_embeddings is not None:
+            scipy.sparse.save_npz(f"{output_dir}/word_embeddings.npz", word_embeddings)
+
+        if train_labels:
+            np.savetxt(f"{output_dir}/train_labels.txt", train_labels, fmt='%i')
+        
+        if train_times is not None:
+            np.savetxt(f"{output_dir}/train_times.txt", train_times, fmt='%i')
+
+        if test_bow is not None:
+            scipy.sparse.save_npz(f"{output_dir}/test_bow.npz", scipy.sparse.csr_matrix(test_bow))
+
+        if test_texts is not None:
+            file_utils.save_text(test_texts, f"{output_dir}/test_texts.txt")
+
+            if test_labels:
+                np.savetxt(f"{output_dir}/test_labels.txt", test_labels, fmt='%i')
+
+            if test_times is not None:
+                np.savetxt(f"{output_dir}/test_times.txt", test_times, fmt='%i')
+        
+        # Save time2id mapping if provided
+        if time2id is not None:
+            with open(f"{output_dir}/time2id.txt", "w", encoding="utf-8") as f:
+                json.dump(time2id, f, indent=2)
+

backend/datasets/data/download.py

@@ -0,0 +1,32 @@
+import os
+import zipfile
+from torchvision.datasets.utils import download_url
+from backend.datasets.utils.logger import Logger
+
+
+logger = Logger("WARNING")
+
+
+def download_dataset(dataset_name, cache_path="~/.topmost"):
+    cache_path = os.path.expanduser(cache_path)
+    raw_filename = f'{dataset_name}.zip'
+
+    if dataset_name in ['Wikitext-103']:
+        # download from Git LFS.
+        zipped_dataset_url = f"https://media.githubusercontent.com/media/BobXWu/TopMost/main/data/{raw_filename}"
+    else:
+        zipped_dataset_url = f"https://raw.githubusercontent.com/BobXWu/TopMost/master/data/{raw_filename}"
+
+    logger.info(zipped_dataset_url)
+
+    download_url(zipped_dataset_url, root=cache_path, filename=raw_filename, md5=None)
+
+    path = f'{cache_path}/{raw_filename}'
+    with zipfile.ZipFile(path, 'r') as zip_ref:
+        zip_ref.extractall(cache_path)
+
+    os.remove(path)
+
+
+if __name__ == '__main__':
+    download_dataset('20NG')

backend/datasets/data/file_utils.py

@@ -0,0 +1,39 @@
+import os
+import json
+
+
+def make_dir(path):
+    os.makedirs(path, exist_ok=True)
+
+
+def read_text(path):
+    texts = list()
+    with open(path, 'r', encoding='utf-8', errors='ignore') as file:
+        for line in file:
+            texts.append(line.strip())
+    return texts
+
+
+def save_text(texts, path):
+    with open(path, 'w', encoding='utf-8') as file:
+        for text in texts:
+            file.write(text.strip() + '\n')
+
+
+def read_jsonlist(path):
+    data = list()
+    with open(path, 'r', encoding='utf-8') as input_file:
+        for line in input_file:
+            data.append(json.loads(line))
+    return data
+
+
+def save_jsonlist(list_of_json_objects, path, sort_keys=True):
+    with open(path, 'w', encoding='utf-8') as output_file:
+        for obj in list_of_json_objects:
+            output_file.write(json.dumps(obj, sort_keys=sort_keys) + '\n')
+
+
+def split_text_word(texts):
+    texts = [text.split() for text in texts]
+    return texts

backend/datasets/dynamic_dataset.py

@@ -0,0 +1,90 @@
+import torch
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+import scipy.sparse
+import scipy.io
+from backend.datasets.data import file_utils
+
+
+class _SequentialDataset(Dataset):
+    def __init__(self, bow, times, time_wordfreq):
+        super().__init__()
+        self.bow = bow
+        self.times = times
+        self.time_wordfreq = time_wordfreq
+
+    def __len__(self):
+        return len(self.bow)
+
+    def __getitem__(self, index):
+        return_dict = {
+            'bow': self.bow[index],
+            'times': self.times[index],
+            'time_wordfreq': self.time_wordfreq[self.times[index]],
+        }
+
+        return return_dict
+
+
+class DynamicDataset:
+    def __init__(self, dataset_dir, batch_size=200, read_labels=False, use_partition=False, device='cuda', as_tensor=True):
+
+        self.load_data(dataset_dir, read_labels, use_partition)
+
+        self.vocab_size = len(self.vocab)
+        self.train_size = len(self.train_bow)
+        self.num_times = int(self.train_times.max()) + 1  # assuming train_times is a numpy array
+        self.train_time_wordfreq = self.get_time_wordfreq(self.train_bow, self.train_times)
+
+        print('train size: ', len(self.train_bow))
+        if use_partition:
+            print('test size: ', len(self.test_bow))
+        print('vocab size: ', len(self.vocab))
+        print('average length: {:.3f}'.format(self.train_bow.sum(1).mean().item()))
+        print('num of each time slice: ', self.num_times, np.bincount(self.train_times))
+
+        if as_tensor:
+            self.train_bow = torch.from_numpy(self.train_bow).float().to(device)
+            self.train_times = torch.from_numpy(self.train_times).long().to(device)
+            self.train_time_wordfreq = torch.from_numpy(self.train_time_wordfreq).float().to(device)
+            
+            if use_partition:
+                self.test_bow = torch.from_numpy(self.test_bow).float().to(device)
+                self.test_times = torch.from_numpy(self.test_times).long().to(device)
+
+            self.train_dataset = _SequentialDataset(self.train_bow, self.train_times, self.train_time_wordfreq)
+            
+            if use_partition:
+                self.test_dataset = _SequentialDataset(self.test_bow, self.test_times, self.train_time_wordfreq)
+
+            self.train_dataloader = DataLoader(self.train_dataset, batch_size=batch_size, shuffle=True)
+
+    def load_data(self, path, read_labels, use_partition=False):
+        self.train_bow = scipy.sparse.load_npz(f'{path}/train_bow.npz').toarray().astype('float32')
+        self.train_texts = file_utils.read_text(f'{path}/train_texts.txt')
+        self.train_times = np.loadtxt(f'{path}/train_times.txt').astype('int32')
+        self.vocab = file_utils.read_text(f'{path}/vocab.txt')
+        self.word_embeddings = scipy.sparse.load_npz(f'{path}/word_embeddings.npz').toarray().astype('float32')
+
+        self.pretrained_WE = self.word_embeddings  # preserve compatibility
+
+        if read_labels:
+            self.train_labels = np.loadtxt(f'{path}/train_labels.txt').astype('int32')
+
+        if use_partition:
+            self.test_bow = scipy.sparse.load_npz(f'{path}/test_bow.npz').toarray().astype('float32')
+            self.test_texts = file_utils.read_text(f'{path}/test_texts.txt')
+            self.test_times = np.loadtxt(f'{path}/test_times.txt').astype('int32')
+            if read_labels:
+                self.test_labels = np.loadtxt(f'{path}/test_labels.txt').astype('int32')
+
+    # word frequency at each time slice.
+    def get_time_wordfreq(self, bow, times):
+        train_time_wordfreq = np.zeros((self.num_times, self.vocab_size))
+        for time in range(self.num_times):
+            idx = np.where(times == time)[0]
+            train_time_wordfreq[time] += bow[idx].sum(0)
+        cnt_times = np.bincount(times)
+
+        train_time_wordfreq = train_time_wordfreq / cnt_times[:, np.newaxis]
+        return train_time_wordfreq

backend/datasets/preprocess.py

@@ -0,0 +1,362 @@
+import json
+import os
+import numpy as np
+from collections import OrderedDict
+import tempfile
+import gensim.downloader
+from tqdm import tqdm
+from backend.datasets.utils.logger import Logger
+import scipy.sparse
+from gensim.models.phrases import Phrases, Phraser
+from typing import List, Union
+from octis.preprocessing.preprocessing import Preprocessing
+
+logger = Logger("WARNING")
+
+class Preprocessor:
+    def __init__(self,
+                 docs_jsonl_path: str,
+                 output_folder: str,
+                 use_partition: bool = False,
+                 use_bigrams: bool = False,
+                 min_count_bigram: int = 5,
+                 threshold_bigram: int = 10,
+                 remove_punctuation: bool = True,
+                 lemmatize: bool = True,
+                 stopword_list: Union[str, List[str]] = None,
+                 min_chars: int = 3,
+                 min_words_docs: int = 10,
+                 min_df: Union[int, float] = 0.0,
+                 max_df: Union[int, float] = 1.0,
+                 max_features: int = None,
+                 language: str = 'english'):
+        
+        self.docs_jsonl_path = docs_jsonl_path
+        self.output_folder = output_folder
+        self.use_partition = use_partition
+        self.use_bigrams = use_bigrams
+        self.min_count_bigram = min_count_bigram
+        self.threshold_bigram = threshold_bigram
+
+        os.makedirs(self.output_folder, exist_ok=True)
+
+        self.preprocessing_params = {
+            'remove_punctuation': remove_punctuation,
+            'lemmatize': lemmatize,
+            'stopword_list': stopword_list,
+            'min_chars': min_chars,
+            'min_words_docs': min_words_docs,
+            'min_df': min_df,
+            'max_df': max_df,
+            'max_features': max_features,
+            'language': language
+        }
+        self.preprocessor_octis = Preprocessing(**self.preprocessing_params)
+
+    def _load_data_to_temp_files(self):
+        """Loads data from JSONL and writes to temporary files for OCTIS preprocessor."""
+        raw_texts = []
+        raw_timestamps = []
+        raw_labels = []
+        has_labels = False
+
+        with open(self.docs_jsonl_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                data = json.loads(line.strip())
+                # Remove newlines from text
+                clean_text = data.get('text', '').replace('\n', ' ').replace('\r', ' ')
+                clean_text = " ".join(clean_text.split())
+                raw_texts.append(clean_text)
+                raw_timestamps.append(data.get('timestamp', ''))
+                label = data.get('label', '')
+                if label:
+                    has_labels = True
+                raw_labels.append(label)
+        
+        # Create temporary files
+        temp_dir = tempfile.mkdtemp()
+        temp_docs_path = os.path.join(temp_dir, "temp_docs.txt")
+        temp_labels_path = None
+
+        with open(temp_docs_path, 'w', encoding='utf-8') as f_docs:
+            for text in raw_texts:
+                f_docs.write(f"{text}\n")
+        
+        if has_labels:
+            temp_labels_path = os.path.join(temp_dir, "temp_labels.txt")
+            with open(temp_labels_path, 'w', encoding='utf-8') as f_labels:
+                for label in raw_labels:
+                    f_labels.write(f"{label}\n")
+
+        print(f"Loaded {len(raw_texts)} raw documents and created temporary files in {temp_dir}.")
+        return raw_texts, raw_timestamps, raw_labels, temp_docs_path, temp_labels_path, temp_dir
+
+    def _make_word_embeddings(self, vocab):
+        """
+        Generates word embeddings for the given vocabulary using GloVe.
+        For n-grams (e.g., "wordA_wordB", "wordX_wordY_wordZ" for n>=2),
+        the resultant embedding is the sum of the embeddings of its constituent
+        single words (wordA + wordB + ...).
+        """
+        print("Loading GloVe word embeddings...")
+        glove_vectors = gensim.downloader.load('glove-wiki-gigaword-200')
+        
+        # Initialize word_embeddings matrix with zeros.
+        # This ensures that words not found (single or n-gram constituents)
+        # will have a zero vector embedding.
+        word_embeddings = np.zeros((len(vocab), glove_vectors.vectors.shape[1]), dtype=np.float32)
+
+        num_found = 0
+
+        try:
+            # Using a set for key_word_list for O(1) average time complexity lookup
+            key_word_list = set(glove_vectors.index_to_key)
+        except AttributeError: # For older gensim versions
+            key_word_list = set(glove_vectors.index2word)
+
+        print("Generating word embeddings for vocabulary (including n-grams)...")
+        for i, word in enumerate(tqdm(vocab, desc="Processing vocabulary words")):
+            if '_' in word: # Check if it's a potential n-gram (n >= 2)
+                parts = word.split('_')
+                
+                # Check if *all* constituent words are present in GloVe
+                all_parts_in_glove = True
+                for part in parts:
+                    if part not in key_word_list:
+                        all_parts_in_glove = False
+                        break # One part not found, stop checking
+                
+                if all_parts_in_glove:
+                    # If all parts are found, sum their embeddings
+                    resultant_vector = np.zeros(glove_vectors.vectors.shape[1], dtype=np.float32)
+                    for part in parts:
+                        resultant_vector += glove_vectors[part]
+                    
+                    word_embeddings[i] = resultant_vector
+                    num_found += 1
+                # Else: one or more constituent words not found, embedding remains zero
+            else: # It's a single word (n=1)
+                if word in key_word_list:
+                    word_embeddings[i] = glove_vectors[word]
+                    num_found += 1
+                # Else: single word not found, embedding remains zero
+
+        logger.info(f'Number of found embeddings (including n-grams): {num_found}/{len(vocab)}')
+        return word_embeddings # Return as dense NumPy array
+    
+    
+    def _save_doc_length_stats(self, filepath: str, output_path: str):
+        doc_lengths = []
+        try:
+            with open(filepath, 'r', encoding='utf-8') as f:
+                for line in f:
+                    doc = line.strip()
+                    if doc:
+                        doc_lengths.append(len(doc))
+        except Exception as e:
+            print(f"Error processing '{filepath}': {e}")
+            return
+
+        if not doc_lengths:
+            print(f"No documents found in '{filepath}'.")
+            return
+
+        stats = {
+            "avg_len": float(np.mean(doc_lengths)),
+            "std_len": float(np.std(doc_lengths)),
+            "max_len": int(np.max(doc_lengths)),
+            "min_len": int(np.min(doc_lengths)),
+            "num_docs": int(len(doc_lengths))
+        }
+
+        with open(output_path, 'w', encoding='utf-8') as f:
+            json.dump(stats, f, indent=4)
+        print(f"Saved document length stats to: {output_path}")
+        
+
+    def preprocess(self):
+        print("Loading data and creating temporary files for OCTIS...")
+        _, raw_timestamps, _, temp_docs_path, temp_labels_path, temp_dir = \
+            self._load_data_to_temp_files()
+
+        print("Starting OCTIS pre-processing using file paths and specified parameters...")
+        octis_dataset = self.preprocessor_octis.preprocess_dataset(
+            documents_path=temp_docs_path,
+            labels_path=temp_labels_path
+        )
+        
+        # Clean up temporary files immediately
+        os.remove(temp_docs_path)
+        if temp_labels_path:
+            os.remove(temp_labels_path)
+        os.rmdir(temp_dir)
+        print(f"Temporary files in {temp_dir} cleaned up.")
+        
+        # --- Proxy: Save __original_indexes and then manually load it ---
+        temp_indexes_dir = tempfile.mkdtemp()
+        temp_indexes_file = os.path.join(temp_indexes_dir, "temp_original_indexes.txt")
+        
+        print(f"Saving __original_indexes to {temp_indexes_file}...")
+        octis_dataset._save_document_indexes(temp_indexes_file)
+        
+        # Manually load the indexes from the file
+        original_indexes_after_octis = []
+        with open(temp_indexes_file, 'r') as f_indexes:
+            for line in f_indexes:
+                original_indexes_after_octis.append(int(line.strip())) # Read as int
+        
+        # Clean up the temporary indexes file and its directory
+        os.remove(temp_indexes_file)
+        os.rmdir(temp_indexes_dir) 
+        print("Temporary indexes file cleaned up.")
+        # --- End Proxy ---
+
+        # Get processed data from OCTIS Dataset object
+        processed_corpus_octis_list = octis_dataset.get_corpus() # List of list of tokens
+        processed_labels_octis = octis_dataset.get_labels() # List of labels
+
+        print("Max index in original_indexes_after_octis:", max(original_indexes_after_octis))
+        print("Length of raw_timestamps:", len(raw_timestamps))
+        
+        # Filter timestamps based on documents that survived OCTIS preprocessing
+        filtered_timestamps = [raw_timestamps[i] for i in original_indexes_after_octis]
+
+        print(f"OCTIS preprocessing complete. {len(processed_corpus_octis_list)} documents remaining.")
+        
+        if self.use_bigrams:
+            print("Generating bigrams with Gensim...")
+            phrases = Phrases(processed_corpus_octis_list, min_count=self.min_count_bigram, threshold=self.threshold_bigram)
+            bigram_phraser = Phraser(phrases)
+            bigrammed_corpus_list = [bigram_phraser[doc] for doc in processed_corpus_octis_list]
+            print("Bigram generation complete.")
+        else:
+            print("Skipping bigram generation as 'use_bigrams' is False.")
+            bigrammed_corpus_list = processed_corpus_octis_list # Use the original processed list
+
+
+        # Convert back to list of strings for easier handling if needed later, but keep as list of lists for BOW
+        bigrammed_texts_for_file = [" ".join(doc) for doc in bigrammed_corpus_list]
+        print("Bigram generation complete.")
+
+        # Build Vocabulary from OCTIS output (after bigrams)
+        # We need a flat list of all tokens to build the vocabulary
+        all_tokens = [token for doc in bigrammed_corpus_list for token in doc]
+        vocab = sorted(list(set(all_tokens))) # Sorted unique words form the vocabulary
+        word_to_id = {word: i for i, word in enumerate(vocab)}
+
+        # Create BOW matrix manually
+        print("Creating Bag-of-Words representations...")
+        rows, cols, data = [], [], []
+        for i, doc_tokens in enumerate(bigrammed_corpus_list):
+            doc_word_counts = {}
+            for token in doc_tokens:
+                if token in word_to_id: # Ensure token is in our final vocab
+                    doc_word_counts[word_to_id[token]] = doc_word_counts.get(word_to_id[token], 0) + 1
+            for col_id, count in doc_word_counts.items():
+                rows.append(i)
+                cols.append(col_id)
+                data.append(count)
+        
+        # Shape is (num_documents, vocab_size)
+        bow_matrix = scipy.sparse.csc_matrix((data, (rows, cols)), shape=(len(bigrammed_corpus_list), len(vocab)))
+        print("Bag-of-Words complete.")
+        
+        # Handle partitioning if required
+        if self.use_partition:
+            num_docs = len(bigrammed_corpus_list)
+            train_size = int(0.8 * num_docs)
+            
+            train_texts = bigrammed_texts_for_file[:train_size]
+            train_bow_matrix = bow_matrix[:train_size]
+            train_timestamps = filtered_timestamps[:train_size]
+            train_labels = processed_labels_octis[:train_size] if processed_labels_octis else []
+
+            test_texts = bigrammed_texts_for_file[train_size:]
+            test_bow_matrix = bow_matrix[train_size:]
+            test_timestamps = filtered_timestamps[train_size:]
+            test_labels = processed_labels_octis[train_size:] if processed_labels_octis else []
+
+        else:
+            train_texts = bigrammed_texts_for_file
+            train_bow_matrix = bow_matrix
+            train_timestamps = filtered_timestamps
+            train_labels = processed_labels_octis
+            test_texts = []
+            test_timestamps = []
+            test_labels = []
+
+        # Generate word embeddings using the provided function
+        word_embeddings = self._make_word_embeddings(vocab)
+
+        # Process timestamps to 0, 1, 2...T and create time2id.txt
+        print("Processing timestamps...")
+        unique_timestamps = sorted(list(set(train_timestamps + test_timestamps)))
+        time_to_id = {timestamp: i for i, timestamp in enumerate(unique_timestamps)}
+
+        train_times_ids = [time_to_id[ts] for ts in train_timestamps]
+        test_times_ids = [time_to_id[ts] for ts in test_timestamps] if self.use_partition else []
+        print("Timestamps processed.")
+
+        # Save files
+        print(f"Saving preprocessed files to {self.output_folder}...")
+        
+        # 1. vocab.txt
+        with open(os.path.join(self.output_folder, "vocab.txt"), "w", encoding="utf-8") as f:
+            for word in vocab:
+                f.write(f"{word}\n")
+
+        # 2. train_texts.txt
+        train_text_path = os.path.join(self.output_folder, "train_texts.txt")
+        with open(train_text_path, "w", encoding="utf-8") as f:
+            for doc in train_texts:
+                f.write(f"{doc}\n")
+        
+        # Save document length stats
+        doc_stats_path = os.path.join(self.output_folder, "length_stats.json")
+        self._save_doc_length_stats(train_text_path, doc_stats_path)
+
+        # 3. train_bow.npz
+        scipy.sparse.save_npz(os.path.join(self.output_folder, "train_bow.npz"), train_bow_matrix)
+
+        # 4. word_embeddings.npz
+        sparse_word_embeddings = scipy.sparse.csr_matrix(word_embeddings)
+        scipy.sparse.save_npz(os.path.join(self.output_folder, "word_embeddings.npz"), sparse_word_embeddings)
+        
+        # 5. train_labels.txt (if labels exist)
+        if train_labels: 
+            with open(os.path.join(self.output_folder, "train_labels.txt"), "w", encoding="utf-8") as f:
+                for label in train_labels:
+                    f.write(f"{label}\n")
+
+        # 6. train_times.txt
+        with open(os.path.join(self.output_folder, "train_times.txt"), "w", encoding="utf-8") as f:
+            for time_id in train_times_ids:
+                f.write(f"{time_id}\n")
+
+        # Files for test set (if use_partition=True)
+        if self.use_partition:
+            # 7. test_bow.npz
+            scipy.sparse.save_npz(os.path.join(self.output_folder, "test_bow.npz"), test_bow_matrix)
+
+            # 8. test_texts.txt
+            with open(os.path.join(self.output_folder, "test_texts.txt"), "w", encoding="utf-8") as f:
+                for doc in test_texts:
+                    f.write(f"{doc}\n")
+
+            # 9. test_labels.txt (if labels exist)
+            if test_labels: 
+                with open(os.path.join(self.output_folder, "test_labels.txt"), "w", encoding="utf-8") as f:
+                    for label in test_labels:
+                        f.write(f"{label}\n")
+
+            # 10. test_times.txt
+            with open(os.path.join(self.output_folder, "test_times.txt"), "w", encoding="utf-8") as f:
+                for time_id in test_times_ids:
+                    f.write(f"{time_id}\n")
+        
+        # 11. time2id.txt
+        sorted_time_to_id = OrderedDict(sorted(time_to_id.items(), key=lambda item: item[1]))
+        with open(os.path.join(self.output_folder, "time2id.txt"), "w", encoding="utf-8") as f:
+            json.dump(sorted_time_to_id, f, indent=4) 
+        
+        print("All files saved successfully.")

backend/datasets/utils/_utils.py

@@ -0,0 +1,37 @@
+import numpy as np
+from backend.datasets.data import file_utils
+
+
+def get_top_words(beta, vocab, num_top_words, verbose=False):
+    topic_str_list = list()
+    for i, topic_dist in enumerate(beta):
+        topic_words = np.array(vocab)[np.argsort(topic_dist)][:-(num_top_words + 1):-1]
+        topic_str = ' '.join(topic_words)
+        topic_str_list.append(topic_str)
+        if verbose:
+            print('Topic {}: {}'.format(i, topic_str))
+
+    return topic_str_list
+
+
+def get_stopwords_set(stopwords=[]):
+    from backend.datasets.data.download import download_dataset
+
+    if stopwords == 'English':
+        from gensim.parsing.preprocessing import STOPWORDS as stopwords
+
+    elif stopwords in ['mallet', 'snowball']:
+        download_dataset('stopwords', cache_path='./')
+        path = f'./stopwords/{stopwords}_stopwords.txt'
+        stopwords = file_utils.read_text(path)
+
+    stopword_set = frozenset(stopwords)
+
+    return stopword_set
+
+
+if __name__ == '__main__':
+    print(list(get_stopwords_set('English'))[:10])
+    print(list(get_stopwords_set('mallet'))[:10])
+    print(list(get_stopwords_set('snowball'))[:10])
+    print(list(get_stopwords_set())[:10])

backend/datasets/utils/logger.py

@@ -0,0 +1,29 @@
+import logging
+
+
+class Logger:
+    def __init__(self, level):
+        self.logger = logging.getLogger('TopMost')
+        self.set_level(level)
+        self._add_handler()
+        self.logger.propagate = False
+
+    def info(self, message):
+        self.logger.info(f"{message}")
+
+    def warning(self, message):
+        self.logger.warning(f"WARNING: {message}")
+
+    def set_level(self, level):
+        levels = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]
+        if level in levels:
+            self.logger.setLevel(level)
+
+    def _add_handler(self):
+        sh = logging.StreamHandler()
+        sh.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(message)s'))
+        self.logger.addHandler(sh)
+
+        # Remove duplicate handlers
+        if len(self.logger.handlers) > 1:
+            self.logger.handlers = [self.logger.handlers[0]]

backend/evaluation/CoherenceModel_ttc.py

@@ -0,0 +1,862 @@
+import logging
+import multiprocessing as mp
+from collections import namedtuple
+
+import numpy as np
+
+from gensim import interfaces, matutils
+from gensim import utils
+from gensim.topic_coherence import (
+    segmentation, probability_estimation,
+    direct_confirmation_measure, indirect_confirmation_measure,
+    aggregation,
+)
+from gensim.topic_coherence.probability_estimation import unique_ids_from_segments
+
+# Set up logging for this module
+logger = logging.getLogger(__name__)
+
+# Define sets for categorizing coherence measures based on their probability estimation method
+BOOLEAN_DOCUMENT_BASED = {'u_mass'}
+SLIDING_WINDOW_BASED = {'c_v', 'c_uci', 'c_npmi', 'c_w2v'}
+
+# Create a namedtuple to define the structure of a coherence measure pipeline
+# Each pipeline consists of a segmentation (seg), probability estimation (prob),
+# confirmation measure (conf), and aggregation (aggr) function.
+_make_pipeline = namedtuple('Coherence_Measure', 'seg, prob, conf, aggr')
+
+# Define the supported coherence measures and their respective pipeline components
+COHERENCE_MEASURES = {
+    'u_mass': _make_pipeline(
+        segmentation.s_one_pre,
+        probability_estimation.p_boolean_document,
+        direct_confirmation_measure.log_conditional_probability,
+        aggregation.arithmetic_mean
+    ),
+    'c_v': _make_pipeline(
+        segmentation.s_one_set,
+        probability_estimation.p_boolean_sliding_window,
+        indirect_confirmation_measure.cosine_similarity,
+        aggregation.arithmetic_mean
+    ),
+    'c_w2v': _make_pipeline(
+        segmentation.s_one_set,
+        probability_estimation.p_word2vec,
+        indirect_confirmation_measure.word2vec_similarity,
+        aggregation.arithmetic_mean
+    ),
+    'c_uci': _make_pipeline(
+        segmentation.s_one_one,
+        probability_estimation.p_boolean_sliding_window,
+        direct_confirmation_measure.log_ratio_measure,
+        aggregation.arithmetic_mean
+    ),
+    'c_npmi': _make_pipeline(
+        segmentation.s_one_one,
+        probability_estimation.p_boolean_sliding_window,
+        direct_confirmation_measure.log_ratio_measure,
+        aggregation.arithmetic_mean
+    ),
+}
+
+# Define default sliding window sizes for different coherence measures
+SLIDING_WINDOW_SIZES = {
+    'c_v': 110,
+    'c_w2v': 5,
+    'c_uci': 10,
+    'c_npmi': 10,
+    'u_mass': None # u_mass does not use a sliding window
+}
+
+
+class CoherenceModel_ttc(interfaces.TransformationABC):
+    """Objects of this class allow for building and maintaining a model for topic coherence.
+
+    Examples
+    ---------
+    One way of using this feature is through providing a trained topic model. A dictionary has to be explicitly provided
+    if the model does not contain a dictionary already
+
+    .. sourcecode:: pycon
+
+        >>> from gensim.test.utils import common_corpus, common_dictionary
+        >>> from gensim.models.ldamodel import LdaModel
+        >>> # Assuming CoherenceModel_ttc is imported or defined in the current scope
+        >>> # from your_module import CoherenceModel_ttc # if saved in a file
+        >>>
+        >>> model = LdaModel(common_corpus, 5, common_dictionary)
+        >>>
+        >>> cm = CoherenceModel_ttc(model=model, corpus=common_corpus, coherence='u_mass')
+        >>> coherence = cm.get_coherence()  # get coherence value
+
+    Another way of using this feature is through providing tokenized topics such as:
+
+    .. sourcecode:: pycon
+
+        >>> from gensim.test.utils import common_corpus, common_dictionary
+        >>> # Assuming CoherenceModel_ttc is imported or defined in the current scope
+        >>> # from your_module import CoherenceModel_ttc # if saved in a file
+        >>> topics = [
+        ...     ['human', 'computer', 'system', 'interface'],
+        ...     ['graph', 'minors', 'trees', 'eps']
+        ... ]
+        >>>
+        >>> cm = CoherenceModel_ttc(topics=topics, corpus=common_corpus, dictionary=common_dictionary, coherence='u_mass')
+        >>> coherence = cm.get_coherence()  # get coherence value
+
+    """
+    def __init__(self, model=None, topics=None, texts=None, corpus=None, dictionary=None,
+                 window_size=None, keyed_vectors=None, coherence='c_v', topn=20, processes=-1):
+        """
+        Initializes the CoherenceModel_ttc.
+
+        Parameters
+        ----------
+        model : :class:`~gensim.models.basemodel.BaseTopicModel`, optional
+            Pre-trained topic model. Should be provided if `topics` is not provided.
+            Supports models that implement the `get_topics` method.
+        topics : list of list of str, optional
+            List of tokenized topics. If provided, `dictionary` must also be provided.
+        texts : list of list of str, optional
+            Tokenized texts, needed for coherence models that use sliding window based (e.g., `c_v`, `c_uci`, `c_npmi`).
+        corpus : iterable of list of (int, number), optional
+            Corpus in Bag-of-Words format.
+        dictionary : :class:`~gensim.corpora.dictionary.Dictionary`, optional
+            Gensim dictionary mapping of id word to create corpus.
+            If `model.id2word` is present and `dictionary` is None, `model.id2word` will be used.
+        window_size : int, optional
+            The size of the window to be used for coherence measures using boolean sliding window as their
+            probability estimator. For 'u_mass' this doesn't matter.
+            If None, default window sizes from `SLIDING_WINDOW_SIZES` are used.
+        keyed_vectors : :class:`~gensim.models.keyedvectors.KeyedVectors`, optional
+            Pre-trained word embeddings (e.g., Word2Vec model) for 'c_w2v' coherence.
+        coherence : {'u_mass', 'c_v', 'c_uci', 'c_npmi', 'c_w2v'}, optional
+            Coherence measure to be used.
+            'u_mass' requires `corpus` (or `texts` which will be converted to corpus).
+            'c_v', 'c_uci', 'c_npmi', 'c_w2v' require `texts`.
+        topn : int, optional
+            Integer corresponding to the number of top words to be extracted from each topic. Defaults to 20.
+        processes : int, optional
+            Number of processes to use for probability estimation phase. Any value less than 1 will be interpreted as
+            `num_cpus - 1`. Defaults to -1.
+        """
+        # Ensure either a model or explicit topics are provided
+        if model is None and topics is None:
+            raise ValueError("One of 'model' or 'topics' has to be provided.")
+        # If topics are provided, a dictionary is mandatory to convert tokens to IDs
+        elif topics is not None and dictionary is None:
+            raise ValueError("Dictionary has to be provided if 'topics' are to be used.")
+
+        self.keyed_vectors = keyed_vectors
+        # Ensure a data source (keyed_vectors, texts, or corpus) is provided for coherence calculation
+        if keyed_vectors is None and texts is None and corpus is None:
+            raise ValueError("One of 'texts', 'corpus', or 'keyed_vectors' has to be provided.")
+
+        # Determine the dictionary to use
+        if dictionary is None:
+            # If no explicit dictionary, try to use the model's dictionary
+            if isinstance(model.id2word, utils.FakeDict):
+                # If model's id2word is a FakeDict, it means no proper dictionary is associated
+                raise ValueError(
+                    "The associated dictionary should be provided with the corpus or 'id2word'"
+                    " for topic model should be set as the associated dictionary.")
+            else:
+                self.dictionary = model.id2word
+        else:
+            self.dictionary = dictionary
+
+        # Store coherence type and window size
+        self.coherence = coherence
+        self.window_size = window_size
+        if self.window_size is None:
+            # Use default window size if not specified
+            self.window_size = SLIDING_WINDOW_SIZES[self.coherence]
+        
+        # Store texts and corpus
+        self.texts = texts
+        self.corpus = corpus
+
+        # Validate inputs based on coherence type
+        if coherence in BOOLEAN_DOCUMENT_BASED:
+            # For document-based measures (e.g., u_mass), corpus is preferred
+            if utils.is_corpus(corpus)[0]:
+                self.corpus = corpus
+            elif self.texts is not None:
+                # If texts are provided, convert them to corpus format
+                self.corpus = [self.dictionary.doc2bow(text) for text in self.texts]
+            else:
+                raise ValueError(
+                    "Either 'corpus' with 'dictionary' or 'texts' should "
+                    "be provided for %s coherence." % coherence)
+
+        elif coherence == 'c_w2v' and keyed_vectors is not None:
+            # For c_w2v, keyed_vectors are needed
+            pass
+        elif coherence in SLIDING_WINDOW_BASED:
+            # For sliding window-based measures, texts are required
+            if self.texts is None:
+                raise ValueError("'texts' should be provided for %s coherence." % coherence)
+        else:
+            # Raise error if coherence type is not supported
+            raise ValueError("%s coherence is not currently supported." % coherence)
+
+        self._topn = topn
+        self._model = model
+        self._accumulator = None  # Cached accumulator for probability estimation
+        self._topics = None       # Store topics internally
+        self.topics = topics      # Call the setter to initialize topics and accumulator state
+
+        # Determine the number of processes to use for parallelization
+        self.processes = processes if processes >= 1 else max(1, mp.cpu_count() - 1)
+
+    @classmethod
+    def for_models(cls, models, dictionary, topn=20, **kwargs):
+        """
+        Initialize a CoherenceModel_ttc with estimated probabilities for all of the given models.
+        This method extracts topics from each model and then uses `for_topics`.
+
+        Parameters
+        ----------
+        models : list of :class:`~gensim.models.basemodel.BaseTopicModel`
+            List of models to evaluate coherence of. Each model should implement
+            the `get_topics` method.
+        dictionary : :class:`~gensim.corpora.dictionary.Dictionary`
+            Gensim dictionary mapping of id word.
+        topn : int, optional
+            Integer corresponding to the number of top words to be extracted from each topic. Defaults to 20.
+        kwargs : object
+            Additional arguments passed to the `CoherenceModel_ttc` constructor (e.g., `corpus`, `texts`, `coherence`).
+
+        Returns
+        -------
+        :class:`~gensim.models.coherencemodel.CoherenceModel`
+            CoherenceModel_ttc instance with estimated probabilities for all given models.
+
+        Example
+        -------
+        .. sourcecode:: pycon
+
+            >>> from gensim.test.utils import common_corpus, common_dictionary
+            >>> from gensim.models.ldamodel import LdaModel
+            >>> # from your_module import CoherenceModel_ttc
+            >>>
+            >>> m1 = LdaModel(common_corpus, 3, common_dictionary)
+            >>> m2 = LdaModel(common_corpus, 5, common_dictionary)
+            >>>
+            >>> cm = CoherenceModel_ttc.for_models([m1, m2], common_dictionary, corpus=common_corpus, coherence='u_mass')
+            >>> # To get coherences for each model:
+            >>> # model_coherences = cm.compare_model_topics([
+            >>> #     CoherenceModel_ttc._get_topics_from_model(m1, topn=cm.topn),
+            >>> #     CoherenceModel_ttc._get_topics_from_model(m2, topn=cm.topn)
+            >>> # ])
+        """
+        # Extract top words as lists for each model's topics
+        topics = [cls.top_topics_as_word_lists(model, dictionary, topn) for model in models]
+        kwargs['dictionary'] = dictionary
+        kwargs['topn'] = topn
+        # Use for_topics to initialize the coherence model with these topics
+        return cls.for_topics(topics, **kwargs)
+
+    @staticmethod
+    def top_topics_as_word_lists(model, dictionary, topn=20):
+        """
+        Get `topn` topics from a model as lists of words.
+
+        Parameters
+        ----------
+        model : :class:`~gensim.models.basemodel.BaseTopicModel`
+            Pre-trained topic model.
+        dictionary : :class:`~gensim.corpora.dictionary.Dictionary`
+            Gensim dictionary mapping of id word.
+        topn : int, optional
+            Integer corresponding to the number of top words to be extracted from each topic. Defaults to 20.
+
+        Returns
+        -------
+        list of list of str
+            Top topics in list-of-list-of-words format.
+        """
+        # Ensure id2token mapping exists in the dictionary
+        if not dictionary.id2token:
+            dictionary.id2token = {v: k for k, v in dictionary.token2id.items()}
+
+        str_topics = []
+        for topic_distribution in model.get_topics():
+            # Get the indices of the topN words based on their probabilities
+            bestn_indices = matutils.argsort(topic_distribution, topn=topn, reverse=True)
+            # Convert word IDs back to words using the dictionary
+            best_words = [dictionary.id2token[_id] for _id in bestn_indices]
+            str_topics.append(best_words)
+        return str_topics
+
+    @classmethod
+    def for_topics(cls, topics_as_topn_terms, **kwargs):
+        """
+        Initialize a CoherenceModel_ttc with estimated probabilities for all of the given topics.
+        This is useful when you have raw topics (list of lists of words) and not a Gensim model object.
+
+        Parameters
+        ----------
+        topics_as_topn_terms : list of list of str
+            Each element in the top-level list should be a list of top-N words, one per topic.
+            For example: `[['word1', 'word2'], ['word3', 'word4']]`.
+
+        Returns
+        -------
+        :class:`~gensim.models.coherencemodel.CoherenceModel`
+            CoherenceModel_ttc with estimated probabilities for the given topics.
+        """
+        if not topics_as_topn_terms:
+            raise ValueError("len(topics_as_topn_terms) must be > 0.")
+        if any(len(topic_list) == 0 for topic_list in topics_as_topn_terms):
+            raise ValueError("Found an empty topic listing in `topics_as_topn_terms`.")
+
+        # Determine the maximum 'topn' value among the provided topics
+        # This will be used to initialize the CoherenceModel_ttc correctly for probability estimation
+        actual_topn_in_data = 0
+        for topic_list in topics_as_topn_terms:
+            for topic in topic_list:
+                actual_topn_in_data = max(actual_topn_in_data, len(topic))
+
+        # Use the provided 'topn' from kwargs, or the determined 'actual_topn_in_data',
+        # ensuring it's not greater than the actual data available.
+        # This allows for precomputing probabilities for a wider set of words if needed.
+        topn_for_prob_estimation = min(kwargs.pop('topn', actual_topn_in_data), actual_topn_in_data)
+
+        # Flatten all topics into a single "super topic" for initial probability estimation.
+        # This ensures that all words relevant to *any* topic in the comparison set
+        # are included in the accumulator.
+        super_topic = utils.flatten(topics_as_topn_terms)
+
+        logger.info(
+            "Number of relevant terms for all %d models (or topic sets): %d",
+            len(topics_as_topn_terms), len(super_topic))
+        
+        # Initialize CoherenceModel_ttc with the super topic to pre-estimate probabilities
+        # for all relevant words across all models.
+        # We pass `topics=[super_topic]` and `topn=len(super_topic)` to ensure all words
+        # are considered during the probability estimation phase.
+        cm = CoherenceModel_ttc(topics=[super_topic], topn=len(super_topic), **kwargs)
+        cm.estimate_probabilities() # Perform the actual probability estimation
+        
+        # After estimation, set the 'topn' back to the desired value for coherence calculation.
+        cm.topn = topn_for_prob_estimation 
+        return cm
+
+    def __str__(self):
+        """Returns a string representation of the coherence measure pipeline."""
+        return str(self.measure)
+
+    @property
+    def model(self):
+        """
+        Get the current topic model used by the instance.
+
+        Returns
+        -------
+        :class:`~gensim.models.basemodel.BaseTopicModel`
+            The currently set topic model.
+        """
+        return self._model
+
+    @model.setter
+    def model(self, model):
+        """
+        Set the topic model for the instance. When a new model is set,
+        it triggers an update of the internal topics and checks if the accumulator needs recomputing.
+
+        Parameters
+        ----------
+        model : :class:`~gensim.models.basemodel.BaseTopicModel`
+            The new topic model to set.
+        """
+        self._model = model
+        if model is not None:
+            new_topics = self._get_topics() # Get topics from the new model
+            self._update_accumulator(new_topics) # Check and update accumulator if needed
+            self._topics = new_topics # Store the new topics
+
+    @property
+    def topn(self):
+        """
+        Get the number of top words (`_topn`) used for coherence calculation.
+
+        Returns
+        -------
+        int
+            The number of top words.
+        """
+        return self._topn
+
+    @topn.setter
+    def topn(self, topn):
+        """
+        Set the number of top words (`_topn`) to consider for coherence calculation.
+        If the new `topn` requires more words than currently loaded topics, and a model is available,
+        it will attempt to re-extract topics from the model.
+
+        Parameters
+        ----------
+        topn : int
+            The new number of top words.
+        """
+        # Get the length of the first topic to check current topic length
+        current_topic_length = len(self._topics[0])
+        # Determine if the new 'topn' requires more words than currently available in topics
+        requires_expansion = current_topic_length < topn
+
+        if self.model is not None:
+            self._topn = topn
+            if requires_expansion:
+                # If expansion is needed and a model is available, re-extract topics from the model.
+                # This call to the setter property `self.model = self._model` effectively re-runs
+                # the logic that extracts topics and updates the accumulator based on the new `_topn`.
+                self.model = self._model
+        else:
+            # If no model is available and expansion is required, raise an error
+            if requires_expansion:
+                raise ValueError("Model unavailable and topic sizes are less than topn=%d" % topn)
+            self._topn = topn  # Topics will be truncated by the `topics` getter if needed
+
+    @property
+    def measure(self):
+        """
+        Returns the namedtuple representing the coherence pipeline functions
+        (segmentation, probability estimation, confirmation, aggregation)
+        based on the `self.coherence` type.
+
+        Returns
+        -------
+        namedtuple
+            Pipeline that contains needed functions/method for calculating coherence.
+        """
+        return COHERENCE_MEASURES[self.coherence]
+
+    @property
+    def topics(self):
+        """
+        Get the current topics. If the internally stored topics have more words
+        than `self._topn`, they are truncated to `self._topn` words.
+
+        Returns
+        -------
+        list of list of str
+            Topics as lists of word tokens.
+        """
+        # If the stored topics contain more words than `_topn`, truncate them
+        if len(self._topics[0]) > self._topn:
+            return [topic[:self._topn] for topic in self._topics]
+        else:
+            return self._topics
+
+    @topics.setter
+    def topics(self, topics):
+        """
+        Set the topics for the instance. This method converts topic words to their
+        corresponding dictionary IDs and updates the accumulator state.
+
+        Parameters
+        ----------
+        topics : list of list of str or list of list of int
+            Topics, either as lists of word tokens or lists of word IDs.
+        """
+        if topics is not None:
+            new_topics = []
+            for topic in topics:
+                # Ensure topic elements are converted to dictionary IDs (numpy array for efficiency)
+                topic_token_ids = self._ensure_elements_are_ids(topic)
+                new_topics.append(topic_token_ids)
+
+            if self.model is not None:
+                # Warn if both model and explicit topics are set, as they might be inconsistent
+                logger.warning(
+                    "The currently set model '%s' may be inconsistent with the newly set topics",
+                    self.model)
+        elif self.model is not None:
+            # If topics are None but a model exists, extract topics from the model
+            new_topics = self._get_topics()
+            logger.debug("Setting topics to those of the model: %s", self.model)
+        else:
+            new_topics = None
+
+        # Check if the accumulator needs to be recomputed based on the new topics
+        self._update_accumulator(new_topics)
+        self._topics = new_topics # Store the (ID-converted) topics
+
+    def _ensure_elements_are_ids(self, topic):
+        """
+        Internal helper to ensure that topic elements are converted to dictionary IDs.
+        Handles cases where input topic might be tokens or already IDs.
+
+        Parameters
+        ----------
+        topic : list of str or list of int
+            A single topic, either as a list of word tokens or word IDs.
+
+        Returns
+        -------
+        :class:`numpy.ndarray`
+            A numpy array of word IDs for the topic.
+
+        Raises
+        ------
+        KeyError
+            If a token is not found in the dictionary or an ID is not a valid key in id2token.
+        """
+        try:
+            # Try to convert tokens to IDs. This is the common case if `topic` contains strings.
+            return np.array([self.dictionary.token2id[token] for token in topic if token in self.dictionary.token2id])
+        except KeyError:
+            # If `KeyError` occurs, assume `topic` might already be a list of IDs.
+            # Attempt to convert IDs to tokens and then back to IDs, ensuring they are valid dictionary entries.
+            # This handles cases where `topic` might contain integer IDs that are not present in the dictionary.
+            try:
+                # Convert IDs to tokens (via id2token) and then tokens to IDs (via token2id)
+                # This filters out invalid IDs.
+                return np.array([self.dictionary.token2id[self.dictionary.id2token[_id]]
+                                 for _id in topic if _id in self.dictionary])
+            except KeyError:
+                raise ValueError("Unable to interpret topic as either a list of tokens or a list of valid IDs within the dictionary.")
+
+    def _update_accumulator(self, new_topics):
+        """
+        Internal helper to determine if the cached `_accumulator` (probability statistics)
+        needs to be wiped and recomputed due to changes in topics.
+        """
+        if self._relevant_ids_will_differ(new_topics):
+            logger.debug("Wiping cached accumulator since it does not contain all relevant ids.")
+            self._accumulator = None
+
+    def _relevant_ids_will_differ(self, new_topics):
+        """
+        Internal helper to check if the set of unique word IDs relevant to the new topics
+        is different from the IDs already covered by the current accumulator.
+
+        Parameters
+        ----------
+        new_topics : list of list of int
+            The new set of topics (as word IDs).
+
+        Returns
+        -------
+        bool
+            True if the relevant IDs will differ, False otherwise.
+        """
+        if self._accumulator is None or not self._topics_differ(new_topics):
+            return False
+
+        # Get unique IDs from the segmented new topics
+        new_set = unique_ids_from_segments(self.measure.seg(new_topics))
+        # Check if the current accumulator's relevant IDs are a superset of the new set.
+        # If not, it means the new topics introduce words not covered, so the accumulator needs updating.
+        return not self._accumulator.relevant_ids.issuperset(new_set)
+
+    def _topics_differ(self, new_topics):
+        """
+        Internal helper to check if the new topics are different from the currently stored topics.
+
+        Parameters
+        ----------
+        new_topics : list of list of int
+            The new set of topics (as word IDs).
+
+        Returns
+        -------
+        bool
+            True if topics are different, False otherwise.
+        """
+        # Compare topic arrays using numpy.array_equal for efficient comparison
+        return (new_topics is not None
+                and self._topics is not None
+                and not np.array_equal(new_topics, self._topics))
+
+    def _get_topics(self):
+        """
+        Internal helper function to extract top words (as IDs) from a trained topic model.
+        """
+        return self._get_topics_from_model(self.model, self.topn)
+
+    @staticmethod
+    def _get_topics_from_model(model, topn):
+        """
+        Internal static method to extract top `topn` words (as IDs) from a trained topic model.
+
+        Parameters
+        ----------
+        model : :class:`~gensim.models.basemodel.BaseTopicModel`
+            Pre-trained topic model (must implement `get_topics` method).
+        topn : int
+            Integer corresponding to the number of top words to extract.
+
+        Returns
+        -------
+        list of :class:`numpy.ndarray`
+            A list where each element is a numpy array of word IDs representing a topic's top words.
+
+        Raises
+        ------
+        AttributeError
+            If the provided model does not implement a `get_topics` method.
+        """
+        try:
+            # Iterate over the topic distributions from the model
+            # Use matutils.argsort to get the indices (word IDs) of the top `topn` words
+            return [
+                matutils.argsort(topic, topn=topn, reverse=True) for topic in
+                model.get_topics()
+            ]
+        except AttributeError:
+            raise ValueError(
+                "This topic model is not currently supported. Supported topic models"
+                " should implement the `get_topics` method.")
+
+    def segment_topics(self):
+        """
+        Segments the current topics using the segmentation function defined by the
+        chosen coherence measure (`self.measure.seg`).
+
+        Returns
+        -------
+        list of list of tuple
+            Segmented topics. The structure depends on the segmentation method (e.g., pairs of word IDs).
+        """
+        # Apply the segmentation function from the pipeline to the current topics
+        return self.measure.seg(self.topics)
+
+    def estimate_probabilities(self, segmented_topics=None):
+        """
+        Accumulates word occurrences and co-occurrences from texts or corpus
+        using the optimal probability estimation method for the chosen coherence metric.
+        This operation can be computationally intensive, especially for sliding window methods.
+
+        Parameters
+        ----------
+        segmented_topics : list of list of tuple, optional
+            Segmented topics. If None, `self.segment_topics()` is called internally.
+
+        Returns
+        -------
+        :class:`~gensim.topic_coherence.text_analysis.CorpusAccumulator`
+            An object that holds the accumulated statistics (word frequencies, co-occurrence frequencies).
+        """
+        if segmented_topics is None:
+            segmented_topics = self.segment_topics()
+
+        # Choose the appropriate probability estimation method based on the coherence type
+        if self.coherence in BOOLEAN_DOCUMENT_BASED:
+            self._accumulator = self.measure.prob(self.corpus, segmented_topics)
+        else:
+            kwargs = dict(
+                texts=self.texts, segmented_topics=segmented_topics,
+                dictionary=self.dictionary, window_size=self.window_size,
+                processes=self.processes)
+            if self.coherence == 'c_w2v':
+                kwargs['model'] = self.keyed_vectors # Pass keyed_vectors for word2vec based coherence
+
+            self._accumulator = self.measure.prob(**kwargs)
+
+        return self._accumulator
+
+    def get_coherence_per_topic(self, segmented_topics=None, with_std=False, with_support=False):
+        """
+        Calculates and returns a list of coherence values, one for each topic,
+        based on the pipeline's confirmation measure.
+
+        Parameters
+        ----------
+        segmented_topics : list of list of tuple, optional
+            Segmented topics. If None, `self.segment_topics()` is called internally.
+        with_std : bool, optional
+            If True, also includes the standard deviation across topic segment sets in addition
+            to the mean coherence for each topic. Defaults to False.
+        with_support : bool, optional
+            If True, also includes the "support" (number of pairwise similarity comparisons)
+            used to compute each topic's coherence. Defaults to False.
+
+        Returns
+        -------
+        list of float or list of tuple
+            A sequence of similarity measures for each topic.
+            If `with_std` or `with_support` is True, each element in the list will be a tuple
+            containing the coherence value and the requested additional statistics.
+        """
+        measure = self.measure
+        if segmented_topics is None:
+            segmented_topics = measure.seg(self.topics)
+        
+        # Ensure probabilities are estimated before calculating coherence
+        if self._accumulator is None:
+            self.estimate_probabilities(segmented_topics)
+
+        kwargs = dict(with_std=with_std, with_support=with_support)
+        if self.coherence in BOOLEAN_DOCUMENT_BASED or self.coherence == 'c_w2v':
+            # These coherence types don't require specific additional kwargs for confirmation measure
+            pass
+        elif self.coherence == 'c_v':
+            # Specific kwargs for c_v's confirmation measure (cosine_similarity)
+            kwargs['topics'] = self.topics
+            kwargs['measure'] = 'nlr' # Normalized Log Ratio
+            kwargs['gamma'] = 1
+        else:
+            # For c_uci and c_npmi, 'normalize' parameter is relevant
+            kwargs['normalize'] = (self.coherence == 'c_npmi')
+
+        return measure.conf(segmented_topics, self._accumulator, **kwargs)
+
+    def aggregate_measures(self, topic_coherences):
+        """
+        Aggregates the individual topic coherence measures into a single overall score
+        using the pipeline's aggregation function (`self.measure.aggr`).
+
+        Parameters
+        ----------
+        topic_coherences : list of float
+            List of coherence values for each topic.
+
+        Returns
+        -------
+        float
+            The aggregated coherence value (e.g., arithmetic mean).
+        """
+        # Apply the aggregation function from the pipeline to the list of topic coherences
+        return self.measure.aggr(topic_coherences)
+
+    def get_coherence(self):
+        """
+        Calculates and returns the overall coherence value for the entire set of topics.
+        This is the main entry point for getting a single coherence score.
+
+        Returns
+        -------
+        float
+            The aggregated coherence value.
+        """
+        # First, get coherence values for each individual topic
+        confirmed_measures = self.get_coherence_per_topic()
+        # Then, aggregate these topic-level coherences into a single score
+        return self.aggregate_measures(confirmed_measures)
+
+    def compare_models(self, models):
+        """
+        Compares multiple topic models by their coherence values.
+        It extracts topics from each model and then calls `compare_model_topics`.
+
+        Parameters
+        ----------
+        models : list of :class:`~gensim.models.basemodel.BaseTopicModel`
+            A sequence of topic models to compare.
+
+        Returns
+        -------
+        list of (list of float, float)
+            A sequence where each element is a pair:
+            (list of average topic coherences for the model, overall model coherence).
+        """
+        # Extract topics (as word IDs) for each model using the internal helper
+        model_topics = [self._get_topics_from_model(model, self.topn) for model in models]
+        # Delegate to compare_model_topics for the actual coherence comparison
+        return self.compare_model_topics(model_topics)
+
+    def compare_model_topics(self, model_topics):
+        """
+        Performs coherence evaluation for each set of topics provided in `model_topics`.
+        This method is designed to be efficient by precomputing probabilities once if needed,
+        and then evaluating coherence for each set of topics.
+
+        Parameters
+        ----------
+        model_topics : list of list of list of int
+            A list where each element is itself a list of topics (each topic being a list of word IDs)
+            representing a set of topics (e.g., from a single model).
+
+        Returns
+        -------
+        list of (list of float, float)
+            A sequence where each element is a pair:
+            (list of average topic coherences for the topic set, overall topic set coherence).
+
+        Notes
+        -----
+        This method uses a heuristic of evaluating coherence at various `topn` values (e.g., 20, 15, 10, 5)
+        and averaging the results for robustness, as suggested in some research.
+        """
+        # Store original topics and topn to restore them after comparison
+        orig_topics = self._topics
+        orig_topn = self.topn
+
+        try:
+            # Perform the actual comparison
+            coherences = self._compare_model_topics(model_topics)
+        finally:
+            # Ensure original topics and topn are restored even if an error occurs
+            self.topics = orig_topics
+            self.topn = orig_topn
+
+        return coherences
+
+    def _compare_model_topics(self, model_topics):
+        """
+        Internal helper to get average topic and model coherences across multiple sets of topics.
+
+        Parameters
+        ----------
+        model_topics : list of list of list of int
+            A list where each element is a set of topics (list of lists of word IDs).
+
+        Returns
+        -------
+        list of (list of float, float)
+            A sequence of pairs:
+            (average topic coherences across different `topn` values for each topic,
+             overall model coherence averaged across different `topn` values).
+        """
+        coherences = []
+        # Define a grid of `topn` values to evaluate coherence.
+        # This provides a more robust average coherence value.
+        # It goes from `self.topn` down to `min(self.topn - 1, 4)` in steps of -5.
+        # e.g., if self.topn is 20, grid might be [20, 15, 10, 5].
+        # The `min(self.topn - 1, 4)` ensures at least some lower values are included,
+        # but also prevents trying `topn` values that are too small or negative.
+        last_topn_value = min(self.topn - 1, 4)
+        topn_grid = list(range(self.topn, last_topn_value, -5))
+        if not topn_grid or max(topn_grid) < 1: # Ensure at least one valid topn if range is empty or too small
+            topn_grid = [max(1, min(self.topn, 5))] # Use min of self.topn and 5, ensure at least 1
+
+        for model_num, topics in enumerate(model_topics):
+            # Set the current topics for the instance to the topics of the model being evaluated
+            self.topics = topics
+
+            coherence_at_n = {} # Dictionary to store coherence results for different `topn` values
+            for n in topn_grid:
+                self.topn = n # Set the `topn` for the current evaluation round
+                topic_coherences = self.get_coherence_per_topic()
+
+                # Handle NaN values in topic coherences by imputing with the mean
+                filled_coherences = np.array(topic_coherences, dtype=float)
+                # Check for NaN values and replace them with the mean of non-NaN values.
+                # np.nanmean handles arrays with all NaNs gracefully by returning NaN.
+                if np.any(np.isnan(filled_coherences)):
+                    mean_val = np.nanmean(filled_coherences)
+                    if np.isnan(mean_val): # If all are NaN, mean_val will also be NaN. In this case, replace with 0 or a very small number.
+                        filled_coherences[np.isnan(filled_coherences)] = 0.0 # Or another sensible default
+                    else:
+                        filled_coherences[np.isnan(filled_coherences)] = mean_val
+
+
+                # Store the topic-level coherences and the aggregated (overall) coherence for this `topn`
+                coherence_at_n[n] = (topic_coherences, self.aggregate_measures(filled_coherences))
+
+            # Unpack the stored coherences for different `topn` values
+            all_topic_coherences_at_n, all_avg_coherences_at_n = zip(*coherence_at_n.values())
+            
+            # Calculate the average topic coherence across all `topn` values
+            # np.vstack stacks lists of topic coherences into a 2D array, then mean(0) computes mean for each topic.
+            avg_topic_coherences = np.vstack(all_topic_coherences_at_n).mean(axis=0)
+            
+            # Calculate the overall model coherence by averaging the aggregated coherences from all `topn` values
+            model_coherence = np.mean(all_avg_coherences_at_n)
+            
+            logging.info("Avg coherence for model %d: %.5f" % (model_num, model_coherence))
+            coherences.append((avg_topic_coherences.tolist(), model_coherence)) # Convert numpy array back to list for output
+
+        return coherences

backend/evaluation/eval.py

@@ -0,0 +1,179 @@
+# dynamic_topic_quality.py
+import numpy as np
+import pandas as pd
+from gensim.corpora.dictionary import Dictionary
+from gensim.models.coherencemodel import CoherenceModel
+from backend.evaluation.CoherenceModel_ttc import CoherenceModel_ttc
+from typing import List, Dict
+
+class TopicQualityAssessor:
+    """
+    Calculates various quality metrics for dynamic topic models from in-memory data.
+
+    This class provides methods to compute:
+    - Temporal Topic Coherence (TTC)
+    - Temporal Topic Smoothness (TTS)
+    - Temporal Topic Quality (TTQ)
+    - Yearly Topic Coherence (TC)
+    - Yearly Topic Diversity (TD)
+    - Yearly Topic Quality (TQ)
+    """
+
+    def __init__(self, topics: List[List[List[str]]], train_texts: List[List[str]], topn: int, coherence_type: str):
+        """
+        Initializes the TopicQualityAssessor with data in memory.
+
+        Args:
+            topics (List[List[List[str]]]): A nested list of topics with structure (T, K, W),
+                                           where T is time slices, K is topics, and W is words.
+            train_texts (List[List[str]]): A list of tokenized documents for the reference corpus.
+            topn (int): Number of top words per topic to consider for calculations.
+            coherence_type (str): The type of coherence to calculate (e.g., 'c_npmi', 'c_v').
+        """
+        # 1. Set texts and dictionary
+        self.texts = train_texts
+        self.dictionary = Dictionary(self.texts)
+
+        # 2. Process topics
+        # User provides topics as (T, K, W) -> List[timestamps][topics][words]
+        # Internal representation for temporal evolution is (K, T, W)
+        topics_array_T_K_W = np.array(topics, dtype=object)
+        if topics_array_T_K_W.ndim != 3:
+            raise ValueError(f"Input 'topics' must be a 3-dimensional list/array. Got {topics_array_T_K_W.ndim} dimensions.")
+        self.total_topics = topics_array_T_K_W.transpose(1, 0, 2) # Shape: (K, T, W)
+        
+        # 3. Get dimensions
+        self.K, self.T, _ = self.total_topics.shape
+
+        # 4. Create topic groups for smoothness calculation (pairs of topics over time)
+        groups = []
+        for k in range(self.K):
+            time_pairs = []
+            for t in range(self.T - 1):
+                time_pairs.append([self.total_topics[k, t].tolist(), self.total_topics[k, t+1].tolist()])
+            groups.append(time_pairs)
+        self.group_topics = np.array(groups, dtype=object)
+
+        # 5. Create yearly topics (T, K, W) for TC/TD calculation
+        self.yearly_topics = self.total_topics.transpose(1, 0, 2)
+        
+        # 6. Set parameters
+        self.topn = topn
+        self.coherence_type = coherence_type
+
+    def _compute_coherence(self, topics: List[List[str]]) -> List[float]:
+        cm = CoherenceModel(
+            topics=topics, texts=self.texts, dictionary=self.dictionary,
+            coherence=self.coherence_type, topn=self.topn
+        )
+        return cm.get_coherence_per_topic()
+
+    def _compute_coherence_ttc(self, topics: List[List[str]]) -> List[float]:
+        cm = CoherenceModel_ttc(
+            topics=topics, texts=self.texts, dictionary=self.dictionary,
+            coherence=self.coherence_type, topn=self.topn
+        )
+        return cm.get_coherence_per_topic()
+
+    def _topic_smoothness(self, topics: List[List[str]]) -> float:
+        K = len(topics)
+        if K <= 1:
+            return 1.0 # Or 0.0, depending on definition. A single topic has no other topic to be dissimilar to.
+        scores = []
+        for i, base in enumerate(topics):
+            base_set = set(base[:self.topn])
+            others = [other for j, other in enumerate(topics) if j != i]
+            if not others:
+                return 1.0
+            overlaps = [len(base_set & set(other[:self.topn])) / self.topn for other in others]
+            scores.append(sum(overlaps) / len(overlaps))
+        return float(sum(scores) / K)
+
+    def get_ttq_dataframe(self) -> pd.DataFrame:
+        """Computes and returns a DataFrame with detailed TTQ metrics per topic chain."""
+        all_coh_scores, avg_coh_scores = [], []
+        for k in range(self.K):
+            coh_per_topic = self._compute_coherence_ttc(self.total_topics[k].tolist())
+            all_coh_scores.append(coh_per_topic)
+            avg_coh_scores.append(float(np.mean(coh_per_topic)))
+
+        all_smooth_scores, avg_smooth_scores = [], []
+        for k in range(self.K):
+            pair_scores = [self._topic_smoothness(pair) for pair in self.group_topics[k]]
+            all_smooth_scores.append(pair_scores)
+            avg_smooth_scores.append(float(np.mean(pair_scores)))
+            
+        df = pd.DataFrame({
+            'topic_idx': list(range(self.K)),
+            'temporal_coherence': all_coh_scores,
+            'temporal_smoothness': all_smooth_scores,
+            'avg_temporal_coherence': avg_coh_scores,
+            'avg_temporal_smoothness': avg_smooth_scores
+        })
+        df['ttq_product'] = df['avg_temporal_coherence'] * df['avg_temporal_smoothness']
+        return df
+
+    def get_tq_dataframe(self) -> pd.DataFrame:
+        """Computes and returns a DataFrame with detailed TQ metrics per time slice."""
+        all_coh, avg_coh, div = [], [], []
+        for t in range(self.T):
+            yearly_t_topics = self.yearly_topics[t].tolist()
+            coh_per_topic = self._compute_coherence(yearly_t_topics)
+            all_coh.append(coh_per_topic)
+            avg_coh.append(float(np.mean(coh_per_topic)))
+            div.append(1 - self._topic_smoothness(yearly_t_topics))
+            
+        df = pd.DataFrame({
+            'year': list(range(self.T)),
+            'all_coherence': all_coh,
+            'avg_coherence': avg_coh,
+            'diversity': div
+        })
+        df['tq_product'] = df['avg_coherence'] * df['diversity']
+        return df
+
+    def get_ttc_score(self) -> float:
+        """Calculates the overall Temporal Topic Coherence (TTC)."""
+        ttq_df = self.get_ttq_dataframe()
+        return ttq_df['avg_temporal_coherence'].mean()
+
+    def get_tts_score(self) -> float:
+        """Calculates the overall Temporal Topic Smoothness (TTS)."""
+        ttq_df = self.get_ttq_dataframe()
+        return ttq_df['avg_temporal_smoothness'].mean()
+
+    def get_ttq_score(self) -> float:
+        """Calculates the overall Temporal Topic Quality (TTQ)."""
+        ttq_df = self.get_ttq_dataframe()
+        return ttq_df['ttq_product'].mean()
+
+    def get_tc_score(self) -> float:
+        """Calculates the overall yearly Topic Coherence (TC)."""
+        tq_df = self.get_tq_dataframe()
+        return tq_df['avg_coherence'].mean()
+
+    def get_td_score(self) -> float:
+        """Calculates the overall yearly Topic Diversity (TD)."""
+        tq_df = self.get_tq_dataframe()
+        return tq_df['diversity'].mean()
+
+    def get_tq_score(self) -> float:
+        """Calculates the overall yearly Topic Quality (TQ)."""
+        tq_df = self.get_tq_dataframe()
+        return tq_df['tq_product'].mean()
+
+    def get_dtq_summary(self) -> Dict[str, float]:
+        """
+        Computes all dynamic topic quality metrics and returns them in a dictionary.
+        """
+        ttq_df = self.get_ttq_dataframe()
+        tq_df = self.get_tq_dataframe()
+        summary = {
+            'TTC': ttq_df['avg_temporal_coherence'].mean(),
+            'TTS': ttq_df['avg_temporal_smoothness'].mean(),
+            'TTQ': ttq_df['ttq_product'].mean(),
+            'TC': tq_df['avg_coherence'].mean(),
+            'TD': tq_df['diversity'].mean(),
+            'TQ': tq_df['tq_product'].mean()
+        }
+        return summary

backend/inference/doc_retriever.py

@@ -0,0 +1,219 @@
+import html
+import json
+import re
+import os
+from hashlib import md5
+
+def deduplicate_docs(collected_docs):
+    seen = set()
+    unique_docs = []
+    for doc in collected_docs:
+        # Prefer unique ID if available
+        key = doc.get("id", md5(doc["text"].encode()).hexdigest())
+        if key not in seen:
+            seen.add(key)
+            unique_docs.append(doc)
+    return unique_docs
+
+def load_length_stats(length_stats_path):
+    """
+    Loads length statistics from a JSON file for a given model path.
+
+    Args:
+        path (str): Path to the model directory containing 'length_stats.json'.
+
+    Returns:
+        dict: A dictionary containing document length statistics.
+    """   
+    if not os.path.exists(length_stats_path):
+        raise FileNotFoundError(f"'length_stats.json' not found at: {length_stats_path}")
+    
+    with open(length_stats_path, "r") as f:
+        length_stats = json.load(f)
+    
+    return length_stats
+
+def get_yearly_counts_for_word(index, word):
+    if word not in index:
+        print(f"[ERROR] Word '{word}' not found in index.")
+        return [], []
+
+    year_counts = index[word]
+    sorted_items = sorted((int(year), len(doc_ids)) for year, doc_ids in year_counts.items())
+    years, counts = zip(*sorted_items) if sorted_items else ([], [])
+    return list(years), list(counts)
+
+
+def get_all_documents_for_word_year(index, docs_file_path, word, year):
+    """
+    Returns all full documents (text + metadata) that contain a given word in a given year.
+
+    Parameters:
+        index (dict): Inverted index.
+        docs_file_path (str): Path to original jsonl corpus.
+        word (str): Word (unigram or bigram).
+        year (int): Year to retrieve docs for.
+
+    Returns:
+        List[Dict]: List of documents with 'id', 'timestamp', and 'text'.
+    """
+    year = int(year)
+
+    if word not in index or year not in index[word]:
+        return []
+
+    doc_ids = set(index[word][year])
+    results = []
+
+    try:
+        with open(docs_file_path, 'r', encoding='utf-8') as f:
+            for doc_id, line in enumerate(f):
+                if doc_id in doc_ids:
+                    doc = json.loads(line)
+                    results.append({
+                        "id": doc_id,
+                        "timestamp": doc.get("timestamp", "N/A"),
+                        "text": doc["text"]
+                    })
+    except Exception as e:
+        print(f"[ERROR] Could not load documents: {e}")
+
+    return results
+
+
+def get_documents_with_all_words_for_year(index, docs_path, words, year):
+    doc_sets = []
+    all_doc_occurrences = {}
+
+    for word in words:
+        word_docs = get_all_documents_for_word_year(index, docs_path, word, year)
+        doc_sets.append(set(doc["id"] for doc in word_docs))
+        for doc in word_docs:
+            all_doc_occurrences.setdefault(doc["id"], doc)
+
+    common_doc_ids = set.intersection(*doc_sets) if doc_sets else set()
+    return [all_doc_occurrences[doc_id] for doc_id in common_doc_ids]
+
+
+def get_intersection_doc_counts_by_year(index, docs_path, words, all_years):
+    year_counts = {}
+    for y in all_years:
+        docs = get_documents_with_all_words_for_year(index, docs_path, words, y)
+        year_counts[y] = len(docs)
+    return year_counts
+
+
+def extract_snippet(text, query, window=30):
+    """
+    Return a short snippet around the first occurrence of the query word.
+    """
+    pattern = re.compile(re.escape(query.replace('_', ' ')), re.IGNORECASE)
+    match = pattern.search(text)
+    if not match:
+        return text[:200] + "..."
+
+    start = max(match.start() - window, 0)
+    end = min(match.end() + window, len(text))
+    snippet = text[start:end].strip()
+
+    return f"...{snippet}..."
+
+def highlight(text, query, highlight_color="#FFD54F"):
+    """
+    Highlight all instances of the query term in text using a colored <mark> tag.
+    """
+    escaped_query = re.escape(query.replace('_', ' '))
+    pattern = re.compile(f"({escaped_query})", flags=re.IGNORECASE)
+
+    def replacer(match):
+        matched_text = html.escape(match.group(1))
+        return f"<mark style='background-color:{highlight_color}; color:black;'>{matched_text}</mark>"
+
+    return pattern.sub(replacer, html.escape(text))
+
+def highlight_words(text, query_words, highlight_color="#24F31D", lemma_to_forms=None):
+    """
+    Highlight all surface forms of each query lemma in the text using a colored <mark> tag.
+
+    Args:
+        text (str): The input raw document text.
+        query_words (List[str]): Lemmatized query tokens to highlight.
+        highlight_color (str): Color to use for highlighting.
+        lemma_to_forms (Dict[str, Set[str]]): Maps a lemma to its surface forms.
+    """
+    # Escape HTML special characters first
+    escaped_text = html.escape(text)
+
+    # Expand query words to include all surface forms
+    expanded_forms = set()
+    for lemma in query_words:
+        if lemma_to_forms and lemma in lemma_to_forms:
+            expanded_forms.update(lemma_to_forms[lemma])
+        else:
+            expanded_forms.add(lemma)  # Fallback if map is missing
+
+    # Sort by length to avoid partial overlaps (e.g., "run" before "running")
+    sorted_queries = sorted(expanded_forms, key=lambda w: -len(w))
+
+    for word in sorted_queries:
+        # Match full word, case insensitive
+        pattern = re.compile(rf'\b({re.escape(word)})\b', flags=re.IGNORECASE)
+
+        def replacer(match):
+            matched_text = match.group(1)
+            return f"<mark style='background-color:{highlight_color}; color:black;'>{matched_text}</mark>"
+
+        escaped_text = pattern.sub(replacer, escaped_text)
+
+    return escaped_text
+
+def get_docs_by_ids(docs_file_path, doc_ids):
+    """
+    Efficiently retrieves specific documents from a .jsonl file by their line number (ID).
+
+    This function reads the file line-by-line and only parses the lines that match
+    the requested document IDs, avoiding loading the entire file into memory.
+
+    Args:
+        docs_file_path (str): The path to the documents.jsonl file.
+        doc_ids (list or set): A collection of document IDs (0-indexed line numbers) to retrieve.
+
+    Returns:
+        list[dict]: A list of document dictionaries that were found. Each dictionary
+                    is augmented with an 'id' key corresponding to its line number.
+    """
+    # Use a set for efficient O(1) lookups.
+    doc_ids_to_find = set(doc_ids)
+    found_docs = {}
+
+    if not doc_ids_to_find:
+        return []
+
+    try:
+        with open(docs_file_path, 'r', encoding='utf-8') as f:
+            for i, line in enumerate(f):
+                # If the current line number is one we're looking for
+                if i in doc_ids_to_find:
+                    try:
+                        doc = json.loads(line)
+                        # Explicitly add the line number as the 'id'
+                        doc['id'] = i
+                        found_docs[i] = doc
+                        # Optimization: stop reading the file once all docs are found
+                        if len(found_docs) == len(doc_ids_to_find):
+                            break
+                    except json.JSONDecodeError:
+                        # Skip malformed lines but inform the user
+                        print(f"[WARNING] Skipping malformed JSON on line {i+1} in {docs_file_path}")
+                        continue
+
+    except FileNotFoundError:
+        print(f"[ERROR] Document file not found at: {docs_file_path}")
+        return []
+    except Exception as e:
+        print(f"[ERROR] An unexpected error occurred while reading documents: {e}")
+        return []
+
+    # Return the documents in the same order as the original doc_ids list
+    # This ensures consistency for downstream processing.
+    return [found_docs[doc_id] for doc_id in doc_ids if doc_id in found_docs]

backend/inference/indexing_utils.py

@@ -0,0 +1,146 @@
+import json
+import os
+import re
+import spacy
+from collections import defaultdict
+
+# Load spaCy once
+nlp = spacy.load("en_core_web_sm", disable=["parser", "ner"])
+
+def tokenize(text):
+    return re.findall(r"\b\w+\b", text.lower())
+
+def has_bigram(tokens, bigram):
+    parts = bigram.split('_')
+    for i in range(len(tokens) - len(parts) + 1):
+        if tokens[i:i + len(parts)] == parts:
+            return True
+    return False
+
+def build_inverse_lemma_map(docs_file_path, cache_path=None):
+    """
+    Build or load a mapping from lemma -> set of surface forms seen in corpus.
+    If cache_path is provided and exists, loads from it.
+    Else builds from scratch and saves to cache_path.
+    """
+    if cache_path and os.path.exists(cache_path):
+        print(f"[INFO] Loading cached lemma_to_forms from {cache_path}")
+        with open(cache_path, "r", encoding="utf-8") as f:
+            raw_map = json.load(f)
+        return {lemma: set(forms) for lemma, forms in raw_map.items()}
+
+    print(f"[INFO] Building inverse lemma map from {docs_file_path}...")
+    lemma_to_forms = defaultdict(set)
+
+    with open(docs_file_path, 'r', encoding='utf-8') as f:
+        for line in f:
+            doc = json.loads(line)
+            tokens = tokenize(doc['text'])
+            spacy_doc = nlp(" ".join(tokens))
+            for token in spacy_doc:
+                lemma_to_forms[token.lemma_].add(token.text.lower())
+
+    if cache_path:
+        print(f"[INFO] Saving lemma_to_forms to {cache_path}")
+        os.makedirs(os.path.dirname(cache_path), exist_ok=True)
+        with open(cache_path, "w", encoding="utf-8") as f:
+            json.dump({k: list(v) for k, v in lemma_to_forms.items()}, f, indent=2)
+
+    return lemma_to_forms
+
+def build_inverted_index(docs_file_path, vocab_set, lemma_map_path=None):
+    vocab_unigrams = {w for w in vocab_set if '_' not in w}
+    vocab_bigrams = {w for w in vocab_set if '_' in w}
+
+    # Load or build lemma map
+    lemma_to_forms = build_inverse_lemma_map(docs_file_path, cache_path=lemma_map_path)
+
+    index = defaultdict(lambda: defaultdict(list))
+    docs = []
+    global_seen_words = set()
+
+    with open(docs_file_path, 'r', encoding='utf-8') as f:
+        for doc_id, line in enumerate(f):
+            doc = json.loads(line)
+            text = doc['text']
+            timestamp = int(doc['timestamp'])
+            docs.append({"text": text, "timestamp": timestamp})
+
+            tokens = tokenize(text)
+            token_set = set(tokens)
+            seen_words = set()
+
+            # Match all lemma queries using surface forms
+            for lemma in vocab_unigrams:
+                surface_forms = lemma_to_forms.get(lemma, set())
+                if token_set & surface_forms:
+                    index[lemma][timestamp].append(doc_id)
+                    seen_words.add(lemma)
+
+            for bigram in vocab_bigrams:
+                if bigram not in seen_words and has_bigram(tokens, bigram):
+                    index[bigram][timestamp].append(doc_id)
+                    seen_words.add(bigram)
+
+            global_seen_words.update(seen_words)
+
+            if (doc_id + 1) % 500 == 0:
+                missing = vocab_set - global_seen_words
+                print(f"[INFO] After {doc_id+1} docs, {len(missing)} vocab words still not seen.")
+                print("Example missing words:", list(missing)[:5])
+
+    missing_final = vocab_set - global_seen_words
+    if missing_final:
+        print(f"[WARNING] {len(missing_final)} vocab words were never found in any document.")
+        print("Examples:", list(missing_final)[:10])
+
+    return index, docs, lemma_to_forms
+
+def save_index_to_disk(index, index_path):
+    index_clean = {
+        word: {str(ts): doc_ids for ts, doc_ids in ts_dict.items()}
+        for word, ts_dict in index.items()
+    }
+    os.makedirs(os.path.dirname(index_path), exist_ok=True)
+    with open(index_path, "w", encoding='utf-8') as f:
+        json.dump(index_clean, f, ensure_ascii=False)
+
+def load_index_from_disk(index_path):
+    with open(index_path, 'r', encoding='utf-8') as f:
+        raw_index = json.load(f)
+
+    index = defaultdict(lambda: defaultdict(list))
+    for word, ts_dict in raw_index.items():
+        for ts, doc_ids in ts_dict.items():
+            index[word][int(ts)] = doc_ids
+
+    return index
+
+def load_docs(docs_file_path):
+    docs = []
+    with open(docs_file_path, 'r', encoding='utf-8') as f:
+        for line in f:
+            doc = json.loads(line)
+            docs.append({
+                "text": doc["text"],
+                "timestamp": int(doc["timestamp"])
+            })
+    return docs
+
+def load_index(docs_file_path, vocab, index_path=None, lemma_map_path=None):
+    if index_path and os.path.exists(index_path):
+        index = load_index_from_disk(index_path)
+        docs = load_docs(docs_file_path)
+        lemma_to_forms = build_inverse_lemma_map(docs_file_path, cache_path=lemma_map_path)
+        return index, docs, lemma_to_forms
+
+    index, docs, lemma_to_forms = build_inverted_index(
+        docs_file_path,
+        set(vocab),
+        lemma_map_path=lemma_map_path
+    )
+
+    if index_path:
+        save_index_to_disk(index, index_path)
+
+    return index, docs, lemma_to_forms

backend/inference/peak_detector.py

@@ -0,0 +1,18 @@
+import numpy as np
+from scipy.signal import find_peaks
+
+def detect_peaks(trend, prominence=0.001, distance=2):
+    """
+    Detect peaks in a word's trend over time.
+
+    Args:
+        trend: List or np.array of floats (word importance over time)
+        prominence: Required prominence of peaks (tune based on scale)
+        distance: Minimum distance between peaks
+
+    Returns:
+        List of indices (timestamps) where peaks occur
+    """
+    trend = np.array(trend)
+    peaks, _ = find_peaks(trend, prominence=prominence, distance=distance)
+    return peaks.tolist()

backend/inference/process_beta.py

@@ -0,0 +1,33 @@
+import numpy as np
+import json
+
+def load_beta_matrix(beta_path: str, vocab_path: str):
+    """
+    Loads the beta matrix (T x K x V) and vocab list.
+
+    Returns:
+        beta: np.ndarray of shape (T, K, V)
+        vocab: list of words
+    """
+    beta = np.load(beta_path)  # shape: T x K x V
+    with open(vocab_path, 'r') as f:
+        vocab = [line.strip() for line in f.readlines()]
+    return beta, vocab
+
+def get_top_words_at_time(beta, vocab, topic_id, time, top_n):
+    topic_beta = beta[time, topic_id, :]
+    top_indices = topic_beta.argsort()[-top_n:][::-1]
+    return [vocab[i] for i in top_indices]
+
+def get_top_words_over_time(beta, vocab, topic_id, top_n):
+    topic_beta = beta[:, topic_id, :]
+    mean_beta = topic_beta.mean(axis=0)
+    top_indices = mean_beta.argsort()[-top_n:][::-1]
+    return [vocab[i] for i in top_indices]
+
+def load_time_labels(time2id_path):
+    with open(time2id_path, 'r') as f:
+        time2id = json.load(f)
+    # Invert and sort by id
+    id2time = {v: k for k, v in time2id.items()}
+    return [id2time[i] for i in sorted(id2time)]

backend/inference/word_selector.py

@@ -0,0 +1,102 @@
+import numpy as np
+from scipy.special import softmax
+
+def get_interesting_words(beta, vocab, topic_id, top_k_final=10, restrict_to=None):
+    """
+    Suggests interesting words by prioritizing "bursty" or "emerging" terms, 
+    making it effective at capturing important low-probability words.
+
+    This algorithm focuses on the ratio of a word's peak probability to its mean,
+    capturing words that show significant growth or have a sudden moment of high
+    relevance, even if their average probability is low.
+    
+    Parameters:
+    - beta: np.ndarray (T, K, V) - Topic-word distributions for each timestamp.
+    - vocab: list of V words - The vocabulary.
+    - topic_id: int - The ID of the topic to analyze.
+    - top_k_final: int - The number of words to return.
+    - restrict_to: optional list of str - Restricts scoring to a subset of words.
+    
+    Returns:
+    - list of top_k_final interesting words (strings).
+    """
+    T, K, V = beta.shape
+    
+    # --- 1. Detect whether softmax is needed ---
+    row_sums = beta.sum(axis=2)
+    is_prob_dist = np.allclose(row_sums, 1.0, atol=1e-2)
+
+    if not is_prob_dist:
+        print("🔁 Beta is not normalized — applying softmax across words per topic.")
+        beta = softmax(beta / 1e-3, axis=2)
+
+    # --- 2. Now extract normalized topic slice ---
+    topic_beta = beta[:, topic_id, :]        # Shape: (T, V)
+    
+    # Mean and Peak probability within the topic for each word
+    mean_topic = topic_beta.mean(axis=0)     # Shape: (V,)
+    peak_topic = topic_beta.max(axis=0)      # Shape: (V,)
+    
+    # Corpus-wide mean for baseline comparison
+    mean_all = beta.mean(axis=(0, 1))        # Shape: (V,)
+
+    # Epsilon to prevent division by zero for words that never appear
+    epsilon = 1e-9
+
+    # --- 3. Calculate the three core components of the new score ---
+    
+    # a) Burstiness Score: How much a word's peak stands out from its own average.
+    # This is the key to finding "surprising" words.
+    burstiness_score = peak_topic / (mean_topic + epsilon)
+
+    # b) Peak Specificity: How much the word's peak in this topic stands out from
+    # its average presence in the entire corpus.
+    peak_specificity_score = peak_topic / (mean_all + epsilon)
+
+    # c) Uniqueness Score (same as before): Penalizes words active in many topics.
+    active_in_topics = (beta > 1e-5).mean(axis=0)  # Shape: (K, V)
+    idf_like = np.log((K + 1) / (active_in_topics.sum(axis=0) + 1)) # Shape: (V,)
+    
+    # --- 4. Compute Final Interestingness Score ---
+    # This score is high for words that are unique, have a high peak relative
+    # to their baseline, and whose peak is an unusual event for that word.
+    final_scores = burstiness_score * peak_specificity_score * idf_like
+    
+    # --- 5. Rank and select top words ---
+    if restrict_to is not None:
+        restrict_set = set(restrict_to)
+        word_indices = [i for i, w in enumerate(vocab) if w in restrict_set]
+    else:
+        word_indices = np.arange(V)
+
+    if not word_indices:
+        return []
+
+    # Rank the filtered indices by the final score in descending order
+    sorted_indices = sorted(word_indices, key=lambda i: -final_scores[i])
+    
+    return [vocab[i] for i in sorted_indices[:top_k_final]]
+
+
+def get_word_trend(beta, vocab, word, topic_id):
+    """
+    Get the time trend of a word's probability under a specific topic.
+
+    Args:
+        beta: np.ndarray of shape (T, K, V)
+        vocab: list of vocab words
+        word: word to search
+        topic_id: index of topic to inspect (0 <= topic_id < K)
+
+    Returns:
+        List of word probabilities over time (length T)
+    """
+    T, K, V = beta.shape
+    if word not in vocab:
+        raise ValueError(f"Word '{word}' not found in vocab.")
+    if not (0 <= topic_id < K):
+        raise ValueError(f"Invalid topic_id {topic_id}. Must be between 0 and {K - 1}.")
+
+    word_index = vocab.index(word)
+    trend = beta[:, topic_id, word_index]  # shape (T,)
+    return trend.tolist()

backend/llm/custom_gemini.py

@@ -0,0 +1,28 @@
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.language_models.chat_models import BaseChatModel
+from typing import List
+
+
+class ChatGemini(BaseChatModel):
+    def __init__(self, api_key: str, model: str = "gemini-pro", temperature: float = 0.7):
+        self.model = model
+        self.temperature = temperature
+        self.api_key = api_key
+        self.client = ChatGoogleGenerativeAI(
+            model=model,
+            temperature=temperature,
+            google_api_key=api_key
+        )
+
+    def _generate(self, messages: List, stop: List[str] = None):
+        # Convert LangChain messages to string
+        prompt = "\n".join(
+            msg.content for msg in messages if isinstance(msg, (HumanMessage, AIMessage))
+        )
+        response = self.client.invoke(prompt)
+        return response
+
+    @property
+    def _llm_type(self) -> str:
+        return "gemini"

backend/llm/custom_mistral.py

@@ -0,0 +1,27 @@
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import HumanMessage, AIMessage
+from langchain_core.outputs import ChatResult, ChatGeneration
+import requests
+import os
+
+class ChatMistral(BaseChatModel):
+    def __init__(self, hf_token=None, model_url=None):
+        self.hf_token = hf_token or os.getenv("HF_TOKEN")
+        self.model_url = model_url or "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.1"
+        self.headers = {"Authorization": f"Bearer {self.hf_token}"}
+
+    def _call(self, prompt: str) -> str:
+        response = requests.post(
+            self.model_url,
+            headers=self.headers,
+            json={"inputs": prompt, "parameters": {"max_new_tokens": 256}},
+        )
+        return response.json()[0]["generated_text"]
+
+    def invoke(self, messages, **kwargs):
+        prompt = "\n".join([msg.content for msg in messages if isinstance(msg, HumanMessage)])
+        response = self._call(prompt)
+        return AIMessage(content=response)
+
+    def _generate(self, messages, stop=None, **kwargs) -> ChatResult:
+        return ChatResult(generations=[ChatGeneration(message=self.invoke(messages))])

backend/llm/llm_router.py

@@ -0,0 +1,73 @@
+from langchain_anthropic import ChatAnthropic
+from backend.llm.custom_mistral import ChatMistral
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_openai import ChatOpenAI
+import os
+import google.auth.transport.requests
+import requests
+
+resp = requests.get("https://www.google.com", proxies={
+    "http": os.getenv("http_proxy"),
+    "https": os.getenv("https_proxy")
+})
+
+def list_supported_models(provider=None):
+    if provider == "OpenAI":
+        return ["gpt-4.1-nano", "gpt-4o-mini"]
+    elif provider == "Anthropic":
+        return ["claude-3-opus-20240229", "claude-3-sonnet-20240229"]
+    elif provider == "Gemini":
+        return ["gemini-2.0-flash-lite", "gemini-1.5-flash"]
+    elif provider == "Mistral":
+        return ["mistral-small", "mistral-medium"]
+    else:
+        # Default fallback: all models grouped by provider
+        return {
+            "OpenAI": ["gpt-4.1-nano", "gpt-4o-mini"],
+            "Anthropic": ["claude-3-opus-20240229", "claude-3-sonnet-20240229"],
+            "Gemini": ["gemini-2.0-flash-lite", "gemini-1.5-flash"],
+            "Mistral": ["mistral-small", "mistral-medium"]
+        }
+
+
+def get_llm(provider: str, model: str, api_key: str = None):
+    if provider == "OpenAI":
+        api_key = api_key or os.getenv("OPENAI_API_KEY")
+        if not api_key:
+            raise ValueError("Missing OpenAI API key.")
+        return ChatOpenAI(model_name=model, temperature=0, openai_api_key=api_key)
+
+    elif provider == "Anthropic":
+        api_key = api_key or os.getenv("ANTHROPIC_API_KEY")
+        if not api_key:
+            raise ValueError("Missing Anthropic API key.")
+        return ChatAnthropic(model=model, temperature=0, anthropic_api_key=api_key)
+
+    elif provider == "Gemini":
+        api_key = api_key or os.getenv("GEMINI_API_KEY")
+        if not api_key:
+            raise ValueError("Missing Gemini API key.")
+        # --- Patch: Set proxy if available ---
+        if "HTTP_PROXY" in os.environ or "http_proxy" in os.environ:
+
+            proxies = {
+                "http": os.getenv("http_proxy") or os.getenv("HTTP_PROXY"),
+                "https": os.getenv("https_proxy") or os.getenv("HTTPS_PROXY")
+            }
+
+            google.auth.transport.requests.requests.Request = lambda *args, **kwargs: requests.Request(
+                *args, **kwargs, proxies=proxies
+            )
+
+        return ChatGoogleGenerativeAI(model=model, temperature=0, google_api_key=api_key)
+
+
+    elif provider == "Mistral":
+        api_key = api_key or os.getenv("MISTRAL_API_KEY")
+        if not api_key:
+            raise ValueError("Missing Mistral API key.")
+        return ChatMistral(model=model, temperature=0, mistral_api_key=api_key)
+
+    else:
+        raise ValueError(f"Unsupported provider: {provider}")
+

backend/llm_utils/label_generator.py

@@ -0,0 +1,72 @@
+from hashlib import sha256
+import json
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from typing import Optional
+import os
+
+#get_top_words_at_time
+from backend.inference.process_beta import get_top_words_at_time
+
+def label_topic_temporal(word_trajectory_str: str, llm, cache_path: Optional[str] = None) -> str:
+    """
+    Label a dynamic topic by providing the LLM with the top words over time.
+
+    Args:
+        word_trajectory_str (str): Formatted keyword evolution string.
+        llm: LangChain-compatible LLM instance.
+        cache_path (Optional[str]): Path to the cache file (JSON).
+
+    Returns:
+        str: Short label for the topic.
+    """
+    topic_key = sha256(word_trajectory_str.encode()).hexdigest()
+
+    # Load cache
+    if cache_path is not None and os.path.exists(cache_path):
+        with open(cache_path, "r") as f:
+            label_cache = json.load(f)
+    else:
+        label_cache = {}
+
+    # Return cached result
+    if topic_key in label_cache:
+        return label_cache[topic_key]
+
+    # Prompt template
+    prompt = ChatPromptTemplate.from_template(
+        "You are an expert in topic modeling and temporal data analysis. "
+        "Given the top words for a topic across multiple time points, your task is to return a short, specific, descriptive topic label. "
+        "Avoid vague, generic, or overly broad labels. Focus on consistent themes in the top words over time. "
+        "Use concise noun phrases, 2–5 words max. Do NOT include any explanation, justification, or extra output.\n\n"
+        "Top words over time:\n{trajectory}\n\n"
+        "Return ONLY the label (no quotes, no extra text):"
+    )
+    chain = prompt | llm | StrOutputParser()
+
+    try:
+        label = chain.invoke({"trajectory": word_trajectory_str}).strip()
+    except Exception as e:
+        label = "Unknown Topic"
+        print(f"[Labeling Error] {e}")
+
+    # Update cache and save
+    label_cache[topic_key] = label
+    if cache_path is not None:
+        os.makedirs(os.path.dirname(cache_path), exist_ok=True)
+        with open(cache_path, "w") as f:
+            json.dump(label_cache, f, indent=2)
+
+    return label
+
+
+def get_topic_labels(beta, vocab, time_labels, llm, cache_path):
+    topic_labels = {}
+    for topic_id in range(beta.shape[1]):
+        word_trajectory_str = "\n".join([
+            f"{time_labels[t]}: {', '.join(get_top_words_at_time(beta, vocab, topic_id, t, top_n=10))}"
+            for t in range(beta.shape[0])
+        ])
+        label = label_topic_temporal(word_trajectory_str, llm=llm, cache_path=cache_path)
+        topic_labels[topic_id] = label
+    return topic_labels

backend/llm_utils/summarizer.py

@@ -0,0 +1,192 @@
+import hashlib
+import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
+from sentence_transformers import SentenceTransformer
+import faiss
+    
+from langchain.prompts import ChatPromptTemplate
+from langchain.docstore.document import Document
+from langchain.memory import ConversationBufferMemory
+from langchain.chains import ConversationChain
+
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+# --- MMR Utilities ---
+def build_mmr_index(docs):
+    texts = [doc['text'] for doc in docs if 'text' in doc]
+    documents = [Document(page_content=text) for text in texts]
+
+    model = SentenceTransformer("all-MiniLM-L6-v2")
+    embeddings = model.encode([doc.page_content for doc in documents], convert_to_numpy=True)
+    faiss.normalize_L2(embeddings)
+
+    index = faiss.IndexFlatIP(embeddings.shape[1])
+    index.add(embeddings)
+
+    return model, index, embeddings, documents
+
+def get_mmr_sample(model, index, embeddings, documents, query, k=15, lambda_mult=0.7):
+    if len(documents) == 0:
+        print("Warning: No documents available, returning empty list.")
+        return []
+
+    if len(documents) <= k:
+        print(f"Warning: Only {len(documents)} documents available, returning all.")
+        return documents
+
+    else:
+        query_vec = model.encode(query, convert_to_numpy=True)
+        query_vec = query_vec / np.linalg.norm(query_vec)
+
+        # Get candidate indices from FAISS (k * 4 or less if not enough documents)
+        num_candidates = min(k * 4, len(documents))
+        D, I = index.search(np.expand_dims(query_vec, axis=0), num_candidates)
+        candidate_idxs = list(I[0])
+
+        selected = []
+        while len(selected) < k and candidate_idxs:
+            if not selected:
+                selected.append(candidate_idxs.pop(0))
+                continue
+
+            mmr_scores = []
+            for idx in candidate_idxs:
+                relevance = cosine_similarity([query_vec], [embeddings[idx]])[0][0]
+                diversity = max([
+                    cosine_similarity([embeddings[idx]], [embeddings[sel]])[0][0]
+                    for sel in selected
+                ])
+                mmr_score = lambda_mult * relevance - (1 - lambda_mult) * diversity
+                mmr_scores.append((idx, mmr_score))
+
+            next_best = max(mmr_scores, key=lambda x: x[1])[0]
+            selected.append(next_best)
+            candidate_idxs.remove(next_best)
+
+        return [documents[i] for i in selected]
+
+
+# --- Summarization ---
+def summarize_docs(word, timestamp, docs, llm, k):
+    if not docs:
+        return "No documents available for this word at this time.", [], 0
+
+    try:
+        model, index, embeddings, documents = build_mmr_index(docs)
+        mmr_docs = get_mmr_sample(model, index, embeddings, documents, query=word, k=k)
+
+        context_texts = "\n".join(f"- {doc.page_content}" for doc in mmr_docs)
+
+        prompt_template = ChatPromptTemplate.from_template(
+            "Given the following documents from {timestamp} containing the word '{word}', "
+            "identify the key themes or distinct discussion points that were prevalent during that time. "
+            "Do NOT describe each bullet in detail. Be concise. Each bullet should be a short phrase or sentence "
+            "capturing a unique, non-overlapping theme. Avoid any elaboration, examples, or justification.\n\n"
+            "Return no more than 5–7 bullets.\n\n"
+            "{context_texts}\n\nSummary:"
+        )
+
+        chain = prompt_template | llm
+        summary = chain.invoke({
+            "word": word,
+            "timestamp": timestamp,
+            "context_texts": context_texts
+        }).content.strip()
+
+        return summary, mmr_docs
+
+    except Exception as e:
+        return f"[Error summarizing: {e}]", [], 0
+
+
+def summarize_multiword_docs(words, timestamp, docs, llm, k):
+    if not docs:
+        return "No common documents available for these words at this time.", []
+
+    try:
+        model, index, embeddings, documents = build_mmr_index(docs)
+        query = " ".join(words)
+        mmr_docs = get_mmr_sample(model, index, embeddings, documents, query=query, k=k)
+
+        context_texts = "\n".join(f"- {doc.page_content}" for doc in mmr_docs)
+
+        prompt_template = ChatPromptTemplate.from_template(
+            "Given the following documents from {timestamp} that all mention the words: '{word_list}', "
+            "identify the key themes or distinct discussion points that were prevalent during that time. "
+            "Do NOT describe each bullet in detail. Be concise. Each bullet should be a short phrase or sentence "
+            "capturing a unique, non-overlapping theme. Avoid any elaboration, examples, or justification.\n\n"
+            "Return no more than 5–7 bullets.\n\n"
+            "{context_texts}\n\n"
+            "Concise Thematic Summary:"
+        )
+
+        chain = prompt_template | llm
+        summary = chain.invoke({
+            "word_list": ", ".join(words),
+            "timestamp": timestamp,
+            "context_texts": context_texts
+        }).content.strip()
+
+        return summary, mmr_docs
+
+    except Exception as e:
+        return f"[Error summarizing: {e}]", []
+
+
+# --- Follow-up Question Handler (Improved) ---
+def ask_multiturn_followup(history: list, question: str, llm, context_texts: str) -> str:
+    """
+    Handles multi-turn follow-up questions based on a provided set of documents.
+
+    This function now REQUIRES context_texts to be provided, ensuring the LLM
+    is always grounded in the source documents for follow-up questions.
+
+    Args:
+        history (list): A list of dictionaries representing the conversation history
+                        (e.g., [{"role": "user", "content": "..."}]).
+        question (str): The user's new follow-up question.
+        llm: The initialized language model instance.
+        context_texts (str): A single string containing all the numbered documents
+                             for context.
+
+    Returns:
+        str: The AI's response to the follow-up question.
+    """
+    try:
+        # 1. Reconstruct conversation memory from the history provided from the UI
+        memory = ConversationBufferMemory(return_messages=True)
+        for turn in history:
+            if turn["role"] == "user":
+                memory.chat_memory.add_user_message(turn["content"])
+            elif turn["role"] == "assistant":
+                memory.chat_memory.add_ai_message(turn["content"])
+
+        # 2. Define the system instruction that grounds the LLM
+        system_instruction = (
+            "You are an assistant answering questions strictly based on the provided sample documents below. "
+            "Your memory contains the previous turns of this conversation. "
+            "If the answer is not clearly available in the text, respond with: "
+            "'The information is not available in the documents provided.'\n\n"
+        )
+
+        # 3. Create the full prompt. No more conditional logic, as context is required.
+        #    The `ConversationChain` will automatically use the memory, so we only need
+        #    to provide the current input, which includes the grounding documents.
+        full_prompt = (
+            f"{system_instruction}"
+            f"--- DOCUMENTS ---\n{context_texts.strip()}\n\n"
+            f"--- QUESTION ---\n{question}"
+        )
+
+        # 4. Create and run the conversation chain
+        conversation = ConversationChain(llm=llm, memory=memory, verbose=False)
+        response = conversation.predict(input=full_prompt)
+        
+        return response.strip()
+
+    except Exception as e:
+        # Good practice to log the full exception for easier debugging
+        print(f"[ERROR] in ask_multiturn_followup: {e}")
+        return f"[Error during multi-turn follow-up. Please check the logs.]"

backend/llm_utils/token_utils.py

@@ -0,0 +1,167 @@
+from typing import Literal
+import tiktoken
+import anthropic
+from typing import List
+
+# Gemini requires the Vertex AI SDK
+try:
+    from vertexai.preview import tokenization as vertex_tokenization
+except ImportError:
+    vertex_tokenization = None
+
+# Mistral requires the SentencePiece tokenizer
+try:
+    import sentencepiece as spm
+except ImportError:
+    spm = None
+
+# ---------------------------
+# Individual Token Counters
+# ---------------------------
+
+def count_tokens_openai(text: str, model_name: str) -> int:
+    try:
+        encoding = tiktoken.encoding_for_model(model_name)
+    except KeyError:
+        encoding = tiktoken.get_encoding("cl100k_base")  # fallback
+    return len(encoding.encode(text))
+
+def count_tokens_anthropic(text: str, model_name: str) -> int:
+    try:
+        client = anthropic.Anthropic()
+        response = client.messages.count_tokens(
+            model=model_name,
+            messages=[{"role": "user", "content": text}]
+        )
+        return response['input_tokens']
+    except Exception as e:
+        raise RuntimeError(f"Anthropic token counting failed: {e}")
+
+def count_tokens_gemini(text: str, model_name: str) -> int:
+    if vertex_tokenization is None:
+        raise ImportError("Please install vertexai: pip install google-cloud-aiplatform[tokenization]")
+    try:
+        tokenizer = vertex_tokenization.get_tokenizer_for_model("gemini-1.5-flash-002")
+        result = tokenizer.count_tokens(text)
+        return result.total_tokens
+    except Exception as e:
+        raise RuntimeError(f"Gemini token counting failed: {e}")
+
+def count_tokens_mistral(text: str) -> int:
+    if spm is None:
+        raise ImportError("Please install sentencepiece: pip install sentencepiece")
+    try:
+        sp = spm.SentencePieceProcessor()
+        # IMPORTANT: You must provide the correct path to the tokenizer model file
+        sp.load("mistral_tokenizer.model")
+        tokens = sp.encode(text, out_type=str)
+        return len(tokens)
+    except Exception as e:
+        raise RuntimeError(f"Mistral token counting failed: {e}")
+
+# ---------------------------
+# Unified Token Counter
+# ---------------------------
+
+def count_tokens(text: str, model_name: str, provider: Literal["OpenAI", "Anthropic", "Gemini", "Mistral"]) -> int:
+    if provider == "OpenAI":
+        return count_tokens_openai(text, model_name)
+    elif provider == "Anthropic":
+        return count_tokens_anthropic(text, model_name)
+    elif provider == "Gemini":
+        return count_tokens_gemini(text, model_name)
+    elif provider == "Mistral":
+        return count_tokens_mistral(text)
+    else:
+        raise ValueError(f"Unsupported provider: {provider}")
+
+
+def get_token_limit_for_model(model_name, provider):
+    # Example values; update as needed for your providers
+    if provider == "openai":
+        if "gpt-4.1-nano" in model_name:
+            return 1047576  # Based on search results
+        elif "gpt-4o-mini" in model_name:
+            return 128000 # Based on search results
+    elif provider == "anthropic":
+        if "claude-3-opus" in model_name:
+            return 200000 # Based on search results
+        elif "claude-3-sonnet" in model_name:
+            return 200000 # Based on search results
+    elif provider == "gemini":
+        if "gemini-2.0-flash-lite" in model_name:
+            return 1048576 # Based on search results
+        elif "gemini-1.5-flash" in model_name:
+            return 1048576 # Based on search results
+    elif provider == "mistral":
+        if "mistral-small" in model_name:
+            return 32000 # Based on search results
+        elif "mistral-medium" in model_name:
+            return 32000 # Based on search results
+    return 8000  # default fallback
+
+
+def estimate_avg_tokens_per_doc(
+    docs: List[str],
+    model_name: str,
+    provider: Literal["OpenAI", "Anthropic", "Gemini", "Mistral"]
+) -> float:
+    """
+    Estimate the average number of tokens per document for the given model.
+
+    Args:
+        docs (List[str]): List of documents.
+        model_name (str): Model name.
+        provider (Literal): LLM provider.
+
+    Returns:
+        float: Average number of tokens per document.
+    """
+    if not docs:
+        return 0.0
+    token_counts = [count_tokens(doc, model_name, provider) for doc in docs]
+    return sum(token_counts) / len(token_counts)
+
+def estimate_max_k(
+    docs: List[str],
+    model_name: str,
+    provider: Literal["OpenAI", "Anthropic", "Gemini", "Mistral"],
+    margin_ratio: float = 0.1,
+) -> int:
+    """
+    Estimate the maximum number of documents that can fit in the context window.
+
+    Returns:
+        int: Estimated K.
+    """
+    if not docs:
+        return 0
+
+    max_tokens = get_token_limit_for_model(model_name, provider)
+    margin = int(max_tokens * margin_ratio)
+    available_tokens = max_tokens - margin
+
+    avg_tokens_per_doc = estimate_avg_tokens_per_doc(docs, model_name, provider)
+    if avg_tokens_per_doc == 0:
+        return 0
+
+    return min(len(docs), int(available_tokens // avg_tokens_per_doc))
+
+def estimate_max_k_fast(docs, margin_ratio=0.1, max_tokens=8000, model_name="gpt-3.5-turbo"):
+    enc = tiktoken.encoding_for_model(model_name)
+    avg_len = sum(len(enc.encode(doc)) for doc in docs[:20]) / min(20, len(docs))
+    margin = int(max_tokens * margin_ratio)
+    available = max_tokens - margin
+    return min(len(docs), int(available // avg_len))
+
+def estimate_k_max_from_word_stats(
+    avg_words_per_doc: float,
+    margin_ratio: float = 0.1,
+    avg_tokens_per_word: float = 1.3,
+    model_name=None,
+    provider=None
+) -> int:
+    model_token_limit = get_token_limit_for_model(model_name, provider)
+    effective_limit = int(model_token_limit * (1 - margin_ratio))
+    est_tokens_per_doc = avg_words_per_doc * avg_tokens_per_word
+    return int(effective_limit // est_tokens_per_doc)

backend/models/CFDTM/CFDTM.py

@@ -0,0 +1,127 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .ETC import ETC
+from .UWE import UWE
+from .Encoder import MLPEncoder
+
+
+class CFDTM(nn.Module):
+    '''
+    Modeling Dynamic Topics in Chain-Free Fashion by Evolution-Tracking Contrastive Learning and Unassociated Word Exclusion. ACL 2024 Findings
+
+    Xiaobao Wu, Xinshuai Dong, Liangming Pan, Thong Nguyen, Anh Tuan Luu.
+    '''
+
+    def __init__(self,
+                 vocab_size,
+                 train_time_wordfreq,
+                 num_times,
+                 pretrained_WE=None,
+                 num_topics=50,
+                 en_units=100,
+                 temperature=0.1,
+                 beta_temp=1.0,
+                 weight_neg=1.0e+7,
+                 weight_pos=1.0e+1,
+                 weight_UWE=1.0e+3,
+                 neg_topk=15,
+                 dropout=0.,
+                 embed_size=200
+                ):
+        super().__init__()
+
+        self.num_topics = num_topics
+        self.beta_temp = beta_temp
+        self.train_time_wordfreq = train_time_wordfreq
+        self.encoder = MLPEncoder(vocab_size, num_topics, en_units, dropout)
+
+        self.a = 1 * np.ones((1, num_topics)).astype(np.float32)
+        self.mu2 = nn.Parameter(torch.as_tensor((np.log(self.a).T - np.mean(np.log(self.a), 1)).T))
+        self.var2 = nn.Parameter(torch.as_tensor((((1.0 / self.a) * (1 - (2.0 / num_topics))).T + (1.0 / (num_topics * num_topics)) * np.sum(1.0 / self.a, 1)).T))
+
+        self.mu2.requires_grad = False
+        self.var2.requires_grad = False
+
+        self.decoder_bn = nn.BatchNorm1d(vocab_size, affine=False)
+
+        if pretrained_WE is None:
+            self.word_embeddings  = nn.init.trunc_normal_(torch.empty(vocab_size, embed_size), std=0.1)
+            self.word_embeddings = nn.Parameter(F.normalize(self.word_embeddings))
+
+        else:
+            self.word_embeddings = nn.Parameter(torch.from_numpy(pretrained_WE).float())
+
+        # topic_embeddings: TxKxD
+        self.topic_embeddings = nn.init.xavier_normal_(torch.zeros(num_topics, self.word_embeddings.shape[1])).repeat(num_times, 1, 1)
+        self.topic_embeddings = nn.Parameter(self.topic_embeddings)
+
+        self.ETC = ETC(num_times, temperature, weight_neg, weight_pos)
+        self.UWE = UWE(self.ETC, num_times, temperature, weight_UWE, neg_topk)
+
+    def get_beta(self):
+        dist = self.pairwise_euclidean_dist(F.normalize(self.topic_embeddings, dim=-1), F.normalize(self.word_embeddings, dim=-1))
+        beta = F.softmax(-dist / self.beta_temp, dim=1)
+
+        return beta
+
+    def pairwise_euclidean_dist(self, x, y):
+        cost = torch.sum(x ** 2, axis=-1, keepdim=True) + torch.sum(y ** 2, axis=-1) - 2 * torch.matmul(x, y.t())
+        return cost
+
+    def get_theta(self, x, times=None):
+        theta, mu, logvar = self.encoder(x)
+        if self.training:
+            return theta, mu, logvar
+
+        return theta
+
+    def get_KL(self, mu, logvar):
+        var = logvar.exp()
+        var_division = var / self.var2
+        diff = mu - self.mu2
+        diff_term = diff * diff / self.var2
+        logvar_division = self.var2.log() - logvar
+        KLD = 0.5 * ((var_division + diff_term + logvar_division).sum(axis=1) - self.num_topics)
+
+        return KLD.mean()
+
+    def get_NLL(self, theta, beta, x, recon_x=None):
+        if recon_x is None:
+            recon_x = self.decode(theta, beta)
+        recon_loss = -(x * recon_x.log()).sum(axis=1)
+
+        return recon_loss
+
+    def decode(self, theta, beta):
+        d1 = F.softmax(self.decoder_bn(torch.bmm(theta.unsqueeze(1), beta).squeeze(1)), dim=-1)
+        return d1
+
+    def forward(self, x, times):
+        loss = 0.
+
+        theta, mu, logvar = self.get_theta(x)
+        kl_theta = self.get_KL(mu, logvar)
+
+        loss += kl_theta
+
+        beta = self.get_beta()
+        time_index_beta = beta[times]
+        recon_x = self.decode(theta, time_index_beta)
+        NLL = self.get_NLL(theta, time_index_beta, x, recon_x)
+        NLL = NLL.mean()
+        loss += NLL
+
+        loss_ETC = self.ETC(self.topic_embeddings)
+        loss += loss_ETC
+
+        loss_UWE = self.UWE(self.train_time_wordfreq, beta, self.topic_embeddings, self.word_embeddings)
+        loss += loss_UWE
+
+        rst_dict = {
+            'loss': loss,
+        }
+
+        return rst_dict

backend/models/CFDTM/ETC.py

@@ -0,0 +1,62 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ETC(nn.Module):
+    def __init__(self, num_times, temperature, weight_neg, weight_pos):
+        super().__init__()
+        self.num_times = num_times
+        self.weight_neg = weight_neg
+        self.weight_pos = weight_pos
+        self.temperature = temperature
+
+    def forward(self, topic_embeddings):
+        loss = 0.
+        loss_neg = 0.
+        loss_pos = 0.
+
+        for t in range(self.num_times):
+            loss_neg += self.compute_loss(topic_embeddings[t], topic_embeddings[t], self.temperature, self_contrast=True)
+
+        for t in range(1, self.num_times):
+            loss_pos += self.compute_loss(topic_embeddings[t], topic_embeddings[t - 1].detach(), self.temperature, self_contrast=False, only_pos=True)
+
+        loss_neg *= (self.weight_neg / self.num_times)
+        loss_pos *= (self.weight_pos / (self.num_times - 1))
+        loss = loss_neg + loss_pos
+
+        return loss
+
+    def compute_loss(self, anchor_feature, contrast_feature, temperature, self_contrast=False, only_pos=False, all_neg=False):
+        # KxK
+        anchor_dot_contrast = torch.div(
+            torch.matmul(F.normalize(anchor_feature, dim=1), F.normalize(contrast_feature, dim=1).T),
+            temperature
+        )
+
+        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
+        logits = anchor_dot_contrast - logits_max.detach()
+
+        pos_mask = torch.eye(anchor_dot_contrast.shape[0]).to(anchor_dot_contrast.device)
+
+        if self_contrast is False:
+            if only_pos is False:
+                if all_neg is True:
+                    exp_logits = torch.exp(logits)
+                    sum_exp_logits = exp_logits.sum(1)
+                    log_prob = -torch.log(sum_exp_logits + 1e-12)
+
+                    mean_log_prob = -log_prob.sum() / (logits.shape[0] * logits.shape[1])
+            else:
+                # only pos
+                mean_log_prob = -(logits * pos_mask).sum() / pos_mask.sum()
+        else:
+            # self contrast: push away from each other in the same time slice.
+            exp_logits = torch.exp(logits) * (1 - pos_mask)
+            sum_exp_logits = exp_logits.sum(1)
+            log_prob = -torch.log(sum_exp_logits + 1e-12)
+
+            mean_log_prob = -log_prob.sum() / (1 - pos_mask).sum()
+
+        return mean_log_prob

backend/models/CFDTM/Encoder.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MLPEncoder(nn.Module):
+    def __init__(self, vocab_size, num_topic, hidden_dim, dropout):
+        super().__init__()
+
+        self.fc11 = nn.Linear(vocab_size, hidden_dim)
+        self.fc12 = nn.Linear(hidden_dim, hidden_dim)
+        self.fc21 = nn.Linear(hidden_dim, num_topic)
+        self.fc22 = nn.Linear(hidden_dim, num_topic)
+
+        self.fc1_drop = nn.Dropout(dropout)
+        self.z_drop = nn.Dropout(dropout)
+
+        self.mean_bn = nn.BatchNorm1d(num_topic, affine=True)
+        self.mean_bn.weight.requires_grad = False
+        self.logvar_bn = nn.BatchNorm1d(num_topic, affine=True)
+        self.logvar_bn.weight.requires_grad = False
+
+    def reparameterize(self, mu, logvar):
+        if self.training:
+            std = torch.exp(0.5 * logvar)
+            eps = torch.randn_like(std)
+            return mu + (eps * std)
+        else:
+            return mu
+
+    def forward(self, x):
+        e1 = F.softplus(self.fc11(x))
+        e1 = F.softplus(self.fc12(e1))
+        e1 = self.fc1_drop(e1)
+        mu = self.mean_bn(self.fc21(e1))
+        logvar = self.logvar_bn(self.fc22(e1))
+        theta = self.reparameterize(mu, logvar)
+        theta = F.softmax(theta, dim=1)
+        theta = self.z_drop(theta)
+        return theta, mu, logvar

backend/models/CFDTM/UWE.py

@@ -0,0 +1,48 @@
+import torch
+import torch.nn as nn
+
+
+class UWE(nn.Module):
+    def __init__(self, ETC, num_times, temperature, weight_UWE, neg_topk):
+        super().__init__()
+
+        self.ETC = ETC
+        self.weight_UWE = weight_UWE
+        self.num_times = num_times
+        self.temperature = temperature
+        self.neg_topk = neg_topk
+
+    def forward(self, time_wordcount, beta, topic_embeddings, word_embeddings):
+        assert(self.num_times == time_wordcount.shape[0])
+
+        topk_indices = self.get_topk_indices(beta)
+
+        loss_UWE = 0.
+        cnt_valid_times = 0.
+        for t in range(self.num_times):
+            neg_idx = torch.where(time_wordcount[t] == 0)[0]
+
+            time_topk_indices = topk_indices[t]
+            neg_idx = list(set(neg_idx.cpu().tolist()).intersection(set(time_topk_indices.cpu().tolist())))
+            neg_idx = torch.tensor(neg_idx).long().to(time_wordcount.device)
+
+            if len(neg_idx) == 0:
+                continue
+
+            time_neg_WE = word_embeddings[neg_idx]
+
+            # topic_embeddings[t]: K x D
+            # word_embeddings[neg_idx]: |V_{neg}| x D
+            loss_UWE += self.ETC.compute_loss(topic_embeddings[t], time_neg_WE, temperature=self.temperature, all_neg=True)
+            cnt_valid_times += 1
+
+        if cnt_valid_times > 0:
+            loss_UWE *= (self.weight_UWE / cnt_valid_times)
+
+        return loss_UWE
+
+    def get_topk_indices(self, beta):
+        # topk_indices: T x K x neg_topk
+        topk_indices = torch.topk(beta, k=self.neg_topk, dim=-1).indices
+        topk_indices = torch.flatten(topk_indices, start_dim=1)
+        return topk_indices

backend/models/DBERTopic_trainer.py

@@ -0,0 +1,99 @@
+import numpy as np
+from bertopic import BERTopic
+from backend.datasets.utils import _utils
+from backend.datasets.utils.logger import Logger
+
+logger = Logger("WARNING")
+
+
+class DBERTopicTrainer:
+    def __init__(self,
+                 dataset,
+                 num_topics=20,
+                 num_top_words=15,
+                 nr_bins=20,
+                 global_tuning=True,
+                 evolution_tuning=True,
+                 datetime_format=None,
+                 verbose=False):
+
+        self.dataset = dataset
+        self.docs = dataset.raw_documents
+        self.num_topics=num_topics
+        # self.timestamps = dataset.train_times
+        self.vocab = dataset.vocab
+        self.num_top_words = num_top_words
+        # self.nr_bins = nr_bins
+        # self.global_tuning = global_tuning
+        # self.evolution_tuning = evolution_tuning
+        # self.datetime_format = datetime_format
+        self.verbose = verbose
+
+        if verbose:
+            logger.set_level("DEBUG")
+        else:
+            logger.set_level("WARNING")
+
+    def train(self, timestamps, datetime_format='%Y'):
+        logger.info("Fitting BERTopic...")
+        self.model = BERTopic(nr_topics=self.num_topics, verbose=self.verbose)
+        self.topics, _ = self.model.fit_transform(self.docs)
+
+        logger.info("Running topics_over_time...")
+        self.topics_over_time_df = self.model.topics_over_time(
+            docs=self.docs,
+            timestamps=timestamps,
+            nr_bins=len(set(timestamps)),
+            datetime_format=datetime_format
+        )
+
+        self.unique_timestamps = sorted(self.topics_over_time_df["Timestamp"].unique())
+        self.unique_topics = sorted(self.topics_over_time_df["Topic"].unique())
+        self.vocab = self.model.vectorizer_model.get_feature_names_out()
+        self.V = len(self.vocab)
+        self.K = len(self.unique_topics)
+        self.T = len(self.unique_timestamps)
+
+    def get_beta(self):
+        logger.info("Generating β matrix...")
+
+        beta = np.zeros((self.T, self.K, self.V))
+        topic_to_index = {topic: idx for idx, topic in enumerate(self.unique_topics)}
+        timestamp_to_index = {timestamp: idx for idx, timestamp in enumerate(self.unique_timestamps)}
+
+        # Extract topic representations at each time
+        for t_idx, timestamp in enumerate(self.unique_timestamps):
+            selection = self.topics_over_time_df[self.topics_over_time_df["Timestamp"] == timestamp]
+            for _, row in selection.iterrows():
+                topic = row["Topic"]
+                words = row["Words"].split(", ")
+                if topic not in topic_to_index:
+                    continue
+                k = topic_to_index[topic]
+                for word in words:
+                    if word in self.vocab:
+                        v = np.where(self.vocab == word)[0][0]
+                        beta[t_idx, k, v] += 1.0
+
+        # Normalize each β_tk to be a probability distribution
+        beta = beta / (beta.sum(axis=2, keepdims=True) + 1e-10)
+        return beta
+
+    def get_top_words(self, num_top_words=None):
+        if num_top_words is None:
+            num_top_words = self.num_top_words
+        beta = self.get_beta()
+        top_words_list = list()
+        for time in range(beta.shape[0]):
+            top_words = _utils.get_top_words(beta[time], self.vocab, num_top_words, self.verbose)
+            top_words_list.append(top_words)
+        return top_words_list
+
+    def get_theta(self):
+        # Not applicable for BERTopic; can return topic assignments or soft topic distributions if required
+        logger.warning("get_theta is not implemented for BERTopic.")
+        return None
+
+    def export_theta(self):
+        logger.warning("export_theta is not implemented for BERTopic.")
+        return None, None

backend/models/DETM.py

@@ -0,0 +1,259 @@
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class DETM(nn.Module):
+    """
+        The Dynamic Embedded Topic Model. 2019
+
+        Adji B. Dieng, Francisco J. R. Ruiz, David M. Blei
+    """
+    def __init__(self, vocab_size, num_times, train_size, train_time_wordfreq,
+                 num_topics=50, train_WE=True, pretrained_WE=None, en_units=800,
+                 eta_hidden_size=200, rho_size=300, enc_drop=0.0, eta_nlayers=3,
+                 eta_dropout=0.0, delta=0.005, theta_act='relu', device='cpu'):
+        super().__init__()
+
+        ## define hyperparameters
+        self.num_topics = num_topics
+        self.num_times = num_times
+        self.vocab_size = vocab_size
+        self.eta_hidden_size = eta_hidden_size
+        self.rho_size = rho_size
+        self.enc_drop = enc_drop
+        self.eta_nlayers = eta_nlayers
+        self.t_drop = nn.Dropout(enc_drop)
+        self.eta_dropout = eta_dropout
+        self.delta = delta
+        self.train_WE = train_WE
+        self.train_size = train_size
+        self.rnn_inp = train_time_wordfreq
+        self.device = device
+
+        self.theta_act = self.get_activation(theta_act)
+
+        ## define the word embedding matrix \rho
+        if self.train_WE:
+            self.rho = nn.Linear(self.rho_size, self.vocab_size, bias=False)
+        else:
+            rho = nn.Embedding(pretrained_WE.size())
+            rho.weight.data = torch.from_numpy(pretrained_WE)
+            self.rho = rho.weight.data.clone().float().to(self.device)
+
+        ## define the variational parameters for the topic embeddings over time (alpha) ... alpha is K x T x L
+        self.mu_q_alpha = nn.Parameter(torch.randn(self.num_topics, self.num_times, self.rho_size))
+        self.logsigma_q_alpha = nn.Parameter(torch.randn(self.num_topics, self.num_times, self.rho_size))
+
+        ## define variational distribution for \theta_{1:D} via amortizartion... theta is K x D
+        self.q_theta = nn.Sequential(
+            nn.Linear(self.vocab_size + self.num_topics, en_units), 
+            self.theta_act,
+            nn.Linear(en_units, en_units),
+            self.theta_act,
+        )
+        self.mu_q_theta = nn.Linear(en_units, self.num_topics, bias=True)
+        self.logsigma_q_theta = nn.Linear(en_units, self.num_topics, bias=True)
+
+        ## define variational distribution for \eta via amortizartion... eta is K x T
+        self.q_eta_map = nn.Linear(self.vocab_size, self.eta_hidden_size)
+        self.q_eta = nn.LSTM(self.eta_hidden_size, self.eta_hidden_size, self.eta_nlayers, dropout=self.eta_dropout)
+        self.mu_q_eta = nn.Linear(self.eta_hidden_size + self.num_topics, self.num_topics, bias=True)
+        self.logsigma_q_eta = nn.Linear(self.eta_hidden_size + self.num_topics, self.num_topics, bias=True)
+
+        self.decoder_bn = nn.BatchNorm1d(vocab_size)
+        self.decoder_bn.weight.requires_grad = False
+
+    def get_activation(self, act):
+        activations = {
+            'tanh': nn.Tanh(),
+            'relu': nn.ReLU(),
+            'softplus': nn.Softplus(),
+            'rrelu': nn.RReLU(),
+            'leakyrelu': nn.LeakyReLU(),
+            'elu': nn.ELU(),
+            'selu': nn.SELU(),
+            'glu': nn.GLU(),
+        }
+
+        if act in activations:
+            act = activations[act]
+        else:
+            print('Defaulting to tanh activations...')
+            act = nn.Tanh()
+        return act
+
+    def reparameterize(self, mu, logvar):
+        """Returns a sample from a Gaussian distribution via reparameterization.
+        """
+        if self.training:
+            std = torch.exp(0.5 * logvar)
+            eps = torch.randn_like(std)
+            return eps.mul_(std).add_(mu)
+        else:
+            return mu
+
+    def get_kl(self, q_mu, q_logsigma, p_mu=None, p_logsigma=None):
+        """Returns KL( N(q_mu, q_logsigma) || N(p_mu, p_logsigma) ).
+        """
+        if p_mu is not None and p_logsigma is not None:
+            sigma_q_sq = torch.exp(q_logsigma)
+            sigma_p_sq = torch.exp(p_logsigma)
+            kl = ( sigma_q_sq + (q_mu - p_mu)**2 ) / ( sigma_p_sq + 1e-6 )
+            kl = kl - 1 + p_logsigma - q_logsigma
+            kl = 0.5 * torch.sum(kl, dim=-1)
+        else:
+            kl = -0.5 * torch.sum(1 + q_logsigma - q_mu.pow(2) - q_logsigma.exp(), dim=-1)
+        return kl
+
+    def get_alpha(self): ## mean field
+        alphas = torch.zeros(self.num_times, self.num_topics, self.rho_size).to(self.device)
+        kl_alpha = []
+
+        alphas[0] = self.reparameterize(self.mu_q_alpha[:, 0, :], self.logsigma_q_alpha[:, 0, :])
+
+        # TODO: why logsigma_p_0 is zero?
+        p_mu_0 = torch.zeros(self.num_topics, self.rho_size).to(self.device)
+        logsigma_p_0 = torch.zeros(self.num_topics, self.rho_size).to(self.device)
+        kl_0 = self.get_kl(self.mu_q_alpha[:, 0, :], self.logsigma_q_alpha[:, 0, :], p_mu_0, logsigma_p_0)
+        kl_alpha.append(kl_0)
+        for t in range(1, self.num_times):
+            alphas[t] = self.reparameterize(self.mu_q_alpha[:, t, :], self.logsigma_q_alpha[:, t, :]) 
+
+            p_mu_t = alphas[t - 1]
+            logsigma_p_t = torch.log(self.delta * torch.ones(self.num_topics, self.rho_size).to(self.device))
+            kl_t = self.get_kl(self.mu_q_alpha[:, t, :], self.logsigma_q_alpha[:, t, :], p_mu_t, logsigma_p_t)
+            kl_alpha.append(kl_t)
+        kl_alpha = torch.stack(kl_alpha).sum()
+        return alphas, kl_alpha.sum()
+
+    def get_eta(self, rnn_inp): ## structured amortized inference
+        inp = self.q_eta_map(rnn_inp).unsqueeze(1)
+        hidden = self.init_hidden()
+        output, _ = self.q_eta(inp, hidden)
+        output = output.squeeze()
+
+        etas = torch.zeros(self.num_times, self.num_topics).to(self.device)
+        kl_eta = []
+
+        inp_0 = torch.cat([output[0], torch.zeros(self.num_topics,).to(self.device)], dim=0)
+        mu_0 = self.mu_q_eta(inp_0)
+        logsigma_0 = self.logsigma_q_eta(inp_0)
+        etas[0] = self.reparameterize(mu_0, logsigma_0)
+
+        p_mu_0 = torch.zeros(self.num_topics,).to(self.device)
+        logsigma_p_0 = torch.zeros(self.num_topics,).to(self.device)
+        kl_0 = self.get_kl(mu_0, logsigma_0, p_mu_0, logsigma_p_0)
+        kl_eta.append(kl_0)
+
+        for t in range(1, self.num_times):
+            inp_t = torch.cat([output[t], etas[t-1]], dim=0)
+            mu_t = self.mu_q_eta(inp_t)
+            logsigma_t = self.logsigma_q_eta(inp_t)
+            etas[t] = self.reparameterize(mu_t, logsigma_t)
+
+            p_mu_t = etas[t-1]
+            logsigma_p_t = torch.log(self.delta * torch.ones(self.num_topics,).to(self.device))
+            kl_t = self.get_kl(mu_t, logsigma_t, p_mu_t, logsigma_p_t)
+            kl_eta.append(kl_t)
+        kl_eta = torch.stack(kl_eta).sum()
+
+        return etas, kl_eta
+
+    def get_theta(self, bows, times, eta=None): ## amortized inference
+        """Returns the topic proportions.
+        """
+
+        normalized_bows = bows / bows.sum(1, keepdims=True)
+
+        if eta is None and self.training is False:
+            eta, kl_eta = self.get_eta(self.rnn_inp)
+
+        eta_td = eta[times]
+        inp = torch.cat([normalized_bows, eta_td], dim=1)
+        q_theta = self.q_theta(inp)
+        if self.enc_drop > 0:
+            q_theta = self.t_drop(q_theta)
+        mu_theta = self.mu_q_theta(q_theta)
+        logsigma_theta = self.logsigma_q_theta(q_theta)
+        z = self.reparameterize(mu_theta, logsigma_theta)
+        theta = F.softmax(z, dim=-1)
+        kl_theta = self.get_kl(mu_theta, logsigma_theta, eta_td, torch.zeros(self.num_topics).to(self.device))
+
+        if self.training:
+            return theta, kl_theta
+        else:
+            return theta
+
+    @property
+    def word_embeddings(self):
+        return self.rho.weight
+
+    @property
+    def topic_embeddings(self):
+        alpha, _ = self.get_alpha()
+        return alpha
+
+    def get_beta(self, alpha=None):
+        """Returns the topic matrix \beta of shape T x K x V
+        """
+
+        if alpha is None and self.training is False:
+            alpha, kl_alpha = self.get_alpha()
+
+        if self.train_WE:
+            logit = self.rho(alpha.view(alpha.size(0) * alpha.size(1), self.rho_size))
+        else:
+            tmp = alpha.view(alpha.size(0) * alpha.size(1), self.rho_size)
+            logit = torch.mm(tmp, self.rho.permute(1, 0))
+        logit = logit.view(alpha.size(0), alpha.size(1), -1)
+
+        beta = F.softmax(logit, dim=-1)
+
+        return beta
+
+    def get_NLL(self, theta, beta, bows):
+        theta = theta.unsqueeze(1)
+        loglik = torch.bmm(theta, beta).squeeze(1)
+        loglik = torch.log(loglik + 1e-12)
+        nll = -loglik * bows
+        nll = nll.sum(-1)
+        return nll
+
+    def forward(self, bows, times):
+        bsz = bows.size(0)
+        coeff = self.train_size / bsz
+        eta, kl_eta = self.get_eta(self.rnn_inp)
+        theta, kl_theta = self.get_theta(bows, times, eta)
+        kl_theta = kl_theta.sum() * coeff
+
+        alpha, kl_alpha = self.get_alpha()
+        beta = self.get_beta(alpha)
+
+        beta = beta[times]
+        # beta = beta[times.type('torch.LongTensor')]
+        nll = self.get_NLL(theta, beta, bows)
+        nll = nll.sum() * coeff
+
+        loss = nll + kl_eta + kl_theta
+
+        rst_dict = {
+            'loss': loss,
+            'nll': nll,
+            'kl_eta': kl_eta,
+            'kl_theta': kl_theta
+        }
+
+        loss += kl_alpha
+        rst_dict['kl_alpha'] = kl_alpha
+
+        return rst_dict
+
+    def init_hidden(self):
+        """Initializes the first hidden state of the RNN used as inference network for \\eta.
+        """
+        weight = next(self.parameters())
+        nlayers = self.eta_nlayers
+        nhid = self.eta_hidden_size
+        return (weight.new_zeros(nlayers, 1, nhid), weight.new_zeros(nlayers, 1, nhid))

backend/models/DTM_trainer.py

@@ -0,0 +1,148 @@
+import gensim
+import numpy as np
+from gensim.models import ldaseqmodel
+from tqdm import tqdm
+import datetime
+from multiprocessing.pool import Pool
+from backend.datasets.utils import _utils
+from backend.datasets.utils.logger import Logger
+
+
+logger = Logger("WARNING")
+
+
+def work(arguments):
+    model, docs = arguments
+    theta_list = list()
+    for doc in tqdm(docs):
+        theta_list.append(model[doc])
+    return theta_list
+
+
+class DTMTrainer:
+    def __init__(self,
+                 dataset,
+                 num_topics=50,
+                 num_top_words=15,
+                 alphas=0.01,
+                 chain_variance=0.005,
+                 passes=10,
+                 lda_inference_max_iter=25,
+                 em_min_iter=6,
+                 em_max_iter=20,
+                 verbose=False
+                ):
+
+        self.dataset = dataset
+        self.vocab_size = dataset.vocab_size
+        self.num_topics = num_topics
+        self.num_top_words = num_top_words
+        self.alphas = alphas
+        self.chain_variance = chain_variance
+        self.passes = passes
+        self.lda_inference_max_iter = lda_inference_max_iter
+        self.em_min_iter = em_min_iter
+        self.em_max_iter = em_max_iter
+
+        self.verbose = verbose
+        if verbose:
+            logger.set_level("DEBUG")
+        else:
+            logger.set_level("WARNING")
+
+    def train(self):
+        id2word = dict(zip(range(self.vocab_size), self.dataset.vocab))
+        train_bow = self.dataset.train_bow
+        train_times = self.dataset.train_times.astype('int32')
+
+        # order documents by time slices
+        self.doc_order_idx = np.argsort(train_times)
+        train_bow = train_bow[self.doc_order_idx]
+        time_slices = np.bincount(train_times)
+
+        corpus = gensim.matutils.Dense2Corpus(train_bow, documents_columns=False)
+
+        self.model = ldaseqmodel.LdaSeqModel(
+            corpus=corpus,
+            id2word=id2word,
+            time_slice=time_slices,
+            num_topics=self.num_topics,
+            alphas=self.alphas,
+            chain_variance=self.chain_variance,
+            em_min_iter=self.em_min_iter,
+            em_max_iter=self.em_max_iter,
+            lda_inference_max_iter=self.lda_inference_max_iter,
+            passes=self.passes
+        )
+
+    def test(self, bow):
+        # bow = dataset.bow.cpu().numpy()
+        # times = dataset.times.cpu().numpy()
+        corpus = gensim.matutils.Dense2Corpus(bow, documents_columns=False)
+
+        num_workers = 20
+        split_idx_list = np.array_split(np.arange(len(bow)), num_workers)
+        worker_size_list = [len(x) for x in split_idx_list]
+
+        worker_id = 0
+        docs_list = [list() for i in range(num_workers)]
+        for i, doc in enumerate(corpus):
+            docs_list[worker_id].append(doc)
+            if len(docs_list[worker_id]) >= worker_size_list[worker_id]:
+                worker_id += 1
+
+        args_list = list()
+        for docs in docs_list:
+            args_list.append([self.model, docs])
+
+        starttime = datetime.datetime.now()
+
+        pool = Pool(processes=num_workers)
+        results = pool.map(work, args_list)
+
+        pool.close()
+        pool.join()
+
+        theta_list = list()
+        for rst in results:
+            theta_list.extend(rst)
+
+        endtime = datetime.datetime.now()
+
+        print("DTM test time: {}s".format((endtime - starttime).seconds))
+
+        return np.asarray(theta_list)
+
+    def get_theta(self):
+        theta = self.model.gammas / self.model.gammas.sum(axis=1)[:, np.newaxis]
+        # NOTE: MUST transform gamma to original order.
+        return theta[np.argsort(self.doc_order_idx)]
+
+    def get_beta(self):
+        beta = list()
+        # K x V x T
+        for item in self.model.topic_chains:
+            # V x T
+            beta.append(item.e_log_prob)
+
+        # T x K x V
+        beta = np.transpose(np.asarray(beta), (2, 0, 1))
+        # use softmax
+        beta = np.exp(beta)
+        beta = beta / beta.sum(-1, keepdims=True)
+        return beta
+
+    def get_top_words(self, num_top_words=None):
+        if num_top_words is None:
+            num_top_words = self.num_top_words
+        beta = self.get_beta()
+        top_words_list = list()
+        for time in range(beta.shape[0]):
+            top_words = _utils.get_top_words(beta[time], self.dataset.vocab, num_top_words, self.verbose)
+            top_words_list.append(top_words)
+        return top_words_list
+
+    def export_theta(self):
+        train_theta = self.get_theta()
+        test_theta = self.test(self.dataset.test_bow)
+        return train_theta, test_theta

backend/models/dynamic_trainer.py

@@ -0,0 +1,177 @@
+import numpy as np
+from tqdm import tqdm
+from collections import defaultdict
+
+import torch
+from torch.optim.lr_scheduler import StepLR
+from backend.datasets.utils import _utils
+from backend.datasets.utils.logger import Logger
+
+logger = Logger("WARNING")
+
+class DynamicTrainer:
+    def __init__(self,
+                 model,
+                 dataset,
+                 num_top_words=15,
+                 epochs=200,
+                 learning_rate=0.002,
+                 batch_size=200,
+                 lr_scheduler=None,
+                 lr_step_size=125,
+                 log_interval=5,
+                 verbose=False
+                ):
+
+        self.model = model
+        self.dataset = dataset
+        self.num_top_words = num_top_words
+        self.epochs = epochs
+        self.learning_rate = learning_rate
+        self.batch_size = batch_size
+        self.lr_scheduler = lr_scheduler
+        self.lr_step_size = lr_step_size
+        self.log_interval = log_interval
+
+        self.verbose = verbose
+        if verbose:
+            logger.set_level("DEBUG")
+        else:
+            logger.set_level("WARNING")
+
+    def make_optimizer(self,):
+        args_dict = {
+            'params': self.model.parameters(),
+            'lr': self.learning_rate,
+        }
+
+        optimizer = torch.optim.Adam(**args_dict)
+        return optimizer
+
+    def make_lr_scheduler(self, optimizer):
+        lr_scheduler = StepLR(optimizer, step_size=self.lr_step_size, gamma=0.5, verbose=False)
+        return lr_scheduler
+
+    def train(self):
+        optimizer = self.make_optimizer()
+
+        if self.lr_scheduler:
+            logger.info("using lr_scheduler")
+            lr_scheduler = self.make_lr_scheduler(optimizer)
+
+        data_size = len(self.dataset.train_dataloader.dataset)
+
+        for epoch in tqdm(range(1, self.epochs + 1)):
+            self.model.train()
+            loss_rst_dict = defaultdict(float)
+
+            for batch_data in self.dataset.train_dataloader:
+
+                rst_dict = self.model(batch_data['bow'], batch_data['times'])
+                batch_loss = rst_dict['loss']
+
+                optimizer.zero_grad()
+                batch_loss.backward()
+                optimizer.step()
+
+                for key in rst_dict:
+                    loss_rst_dict[key] += rst_dict[key] * len(batch_data)
+
+            if self.lr_scheduler:
+                lr_scheduler.step()
+
+            if epoch % self.log_interval == 0:
+                output_log = f'Epoch: {epoch:03d}'
+                for key in loss_rst_dict:
+                    output_log += f' {key}: {loss_rst_dict[key] / data_size :.3f}'
+
+                logger.info(output_log)
+
+        top_words = self.get_top_words()
+        train_theta = self.test(self.dataset.train_bow, self.dataset.train_times)
+
+        return top_words, train_theta
+
+    def test(self, bow, times):
+        data_size = bow.shape[0]
+        theta = list()
+        all_idx = torch.split(torch.arange(data_size), self.batch_size)
+
+        with torch.no_grad():
+            self.model.eval()
+            for idx in all_idx:
+                batch_theta = self.model.get_theta(bow[idx], times[idx])
+                theta.extend(batch_theta.cpu().tolist())
+
+        theta = np.asarray(theta)
+        return theta
+
+    def get_beta(self):
+        self.model.eval()
+        beta = self.model.get_beta().detach().cpu().numpy()
+        return beta
+
+    def get_top_words(self, num_top_words=None):
+        if num_top_words is None:
+            num_top_words = self.num_top_words
+
+        beta = self.get_beta()
+        top_words_list = list()
+        for time in range(beta.shape[0]):
+            if self.verbose:
+                print(f"======= Time: {time} =======")
+            top_words = _utils.get_top_words(beta[time], self.dataset.vocab, num_top_words, self.verbose)
+            top_words_list.append(top_words)
+        return top_words_list
+
+    def export_theta(self):
+        train_theta = self.test(self.dataset.train_bow, self.dataset.train_times)
+        test_theta = self.test(self.dataset.test_bow, self.dataset.test_times)
+
+        return train_theta, test_theta
+    
+    def get_top_words_at_time(self, topic_id, time, top_n):
+        beta = self.get_beta()  # shape: [T, K, V]
+        topic_beta = beta[time, topic_id, :]
+        top_indices = topic_beta.argsort()[-top_n:][::-1]
+        return [self.dataset.vocab[i] for i in top_indices]
+    
+
+    def get_topic_words_over_time(self, topic_id, top_n):
+        """
+        Returns top_n words for the given topic_id over all time steps.
+        Output: List[List[str]], each inner list is the top_n words at a time step.
+        """
+        beta = self.get_beta()  # shape: [T, K, V]
+        T = beta.shape[0]
+        return [
+            self.get_top_words_at_time(topic_id=topic_id, time=t, top_n=top_n)
+            for t in range(T)
+        ]
+
+    def get_all_topics_at_time(self, time, top_n):
+        """
+        Returns top_n words for each topic at the given time step.
+        Output: List[List[str]], each inner list is the top_n words for a topic.
+        """
+        beta = self.get_beta()  # shape: [T, K, V]
+        K = beta.shape[1]
+        return [
+            self.get_top_words_at_time(topic_id=k, time=time, top_n=top_n)
+            for k in range(K)
+        ]
+    
+    def get_all_topics_over_time(self, top_n=10):
+        """
+        Returns the top_n words for all topics over all time steps.
+        Output shape: List[List[List[str]]] = T x K x top_n
+        """
+        beta = self.get_beta()  # shape: [T, K, V]
+        T, K, _ = beta.shape
+        return [
+            [
+                self.get_top_words_at_time(topic_id=k, time=t, top_n=top_n)
+                for k in range(K)
+            ]
+            for t in range(T)
+        ]

data/ACL_Anthology/CFDTM/beta.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:34984bfb432a10733161a9dfed834a9ef4f366a28a6cb2ecd6e8351997f1599a
+size 16645248

data/ACL_Anthology/DETM/beta.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c6eefa9b6aaea4c694736d09ad9e517446f09929c01889e26633300e5eff166
+size 41612928

data/ACL_Anthology/DTM/beta.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:14c296a2e3fb49f9d0b66262907d64f7d181408768e43138d57c262ea6a11318
+size 33290368

data/ACL_Anthology/DTM/topic_label_cache.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ea9f3c508ede82967cdf02050d7383d58dd9d269a7f661ae1462a95cbac3331e
+size 2089

data/ACL_Anthology/docs.jsonl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a004dd095b9a4f29fdccb5144d50d3dacc7985af443a8de434005b7b8401f9b7
+size 67395059

data/ACL_Anthology/inverted_index.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:60e7ee888abb2fd025b11415a7ead6780d41c5f890cc25ba453615906f10b8d7
+size 30865281

data/ACL_Anthology/processed/lemma_to_forms.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:00ea8855f9ced2ca3d785ce5926ced29b35e0779cd6b3166edfd5c5a5c1beccb
+size 4370995

data/ACL_Anthology/processed/length_stats.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5cc985e5a1ce565ca4179d343ade1526daab463520f6317122953da83d368306
+size 133

data/ACL_Anthology/processed/time2id.txt

@@ -0,0 +1,18 @@
+{
+    "2010": 0,
+    "2011": 1,
+    "2012": 2,
+    "2013": 3,
+    "2014": 4,
+    "2015": 5,
+    "2016": 6,
+    "2017": 7,
+    "2018": 8,
+    "2019": 9,
+    "2020": 10,
+    "2021": 11,
+    "2022": 12,
+    "2023": 13,
+    "2024": 14,
+    "2025": 15
+}