| 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 |
|
|
