|
|
import os |
|
|
|
|
|
import spaces |
|
|
from bertopic import BERTopic |
|
|
from bertopic.representation import KeyBERTInspired, MaximalMarginalRelevance |
|
|
from hdbscan import HDBSCAN |
|
|
from sentence_transformers import SentenceTransformer |
|
|
from umap import UMAP |
|
|
|
|
|
from src.utils.constants import EMBEDDING_MODEL_NAME, MODEL_REPO_ID |
|
|
from src.utils.utils import get_timestamp |
|
|
|
|
|
HF_TOKEN = os.environ.get("HF_TOKEN", None) |
|
|
embedding_model = SentenceTransformer(EMBEDDING_MODEL_NAME) |
|
|
|
|
|
|
|
|
@spaces.GPU(duration=120) |
|
|
def topic_modeling( |
|
|
filename, |
|
|
docs, |
|
|
embeddings, |
|
|
n_gram_range=(3, 6), |
|
|
mmr_diversity=1, |
|
|
mmr_top_n_words=30, |
|
|
keybert_top_n_words=50, |
|
|
random_state=42, |
|
|
min_cluster_size=15, |
|
|
): |
|
|
""" |
|
|
Perform topic modeling on a list of documents and their embeddings. |
|
|
|
|
|
Parameters |
|
|
---------- |
|
|
docs : List of str |
|
|
The list of documents to be topic modeled. |
|
|
embeddings : List of numpy.ndarray |
|
|
The list of embeddings of the given documents. |
|
|
embedding_model : SentenceTransformer |
|
|
The embedding model used to generate the embeddings. |
|
|
n_gram_range : Tuple of int, optional |
|
|
The range of n-grams to be considered. Defaults to (3, 6). |
|
|
mmr_diversity : float, optional |
|
|
The diversity value of the MMR model. Defaults to 1. |
|
|
mmr_top_n_words : int, optional |
|
|
The number of top words to be considered in the MMR model. Defaults to 30. |
|
|
keybert_top_n_words : int, optional |
|
|
The number of top words to be considered in the KeyBERT model. Defaults to 50. |
|
|
random_state : int, optional |
|
|
The random seed for reproducibility. Defaults to 42. |
|
|
min_cluster_size : int, optional |
|
|
The minimum size of a cluster to be considered as a topic. Defaults to 15. |
|
|
|
|
|
Returns |
|
|
------- |
|
|
fig : matplotlib.figure.Figure |
|
|
The datamap of the topic modeling. |
|
|
topic_info_df : pandas.DataFrame |
|
|
The topic information dataframe. |
|
|
""" |
|
|
representation_model = [ |
|
|
KeyBERTInspired(top_n_words=keybert_top_n_words, random_state=random_state), |
|
|
MaximalMarginalRelevance(diversity=mmr_diversity, top_n_words=mmr_top_n_words), |
|
|
] |
|
|
|
|
|
hdbscan_model = HDBSCAN( |
|
|
min_cluster_size=min_cluster_size, |
|
|
metric="euclidean", |
|
|
cluster_selection_method="eom", |
|
|
prediction_data=True, |
|
|
) |
|
|
|
|
|
umap_model = UMAP( |
|
|
n_neighbors=15, |
|
|
n_components=5, |
|
|
min_dist=0.0, |
|
|
metric="cosine", |
|
|
low_memory=False, |
|
|
random_state=random_state, |
|
|
) |
|
|
|
|
|
topic_model = BERTopic( |
|
|
embedding_model=embedding_model, |
|
|
representation_model=representation_model, |
|
|
n_gram_range=n_gram_range, |
|
|
hdbscan_model=hdbscan_model, |
|
|
umap_model=umap_model, |
|
|
verbose=True, |
|
|
).fit(docs, embeddings=embeddings) |
|
|
|
|
|
topic_model.push_to_hf_hub( |
|
|
repo_id=MODEL_REPO_ID, |
|
|
commit_message=f"{get_timestamp()} - {filename}", |
|
|
token=HF_TOKEN, |
|
|
private=True, |
|
|
serialization="safetensors", |
|
|
save_embedding_model=EMBEDDING_MODEL_NAME, |
|
|
save_ctfidf=True, |
|
|
) |
|
|
|
|
|
topic_info_df = topic_model.get_topic_info() |
|
|
|
|
|
return topic_info_df |
|
|
|
