topicBuzz / test /bertopicTest.py
Wyomike's picture
Init commit for hf
de99ac3
Raw
History Blame Contribute Delete
1.06 kB
import chromadb
client = chromadb.PersistentClient(path="./my_mastodon_db")
collection = client.get_collection(name="mastodon_posts")
# This might take a moment.
print("Fetching all data from ChromaDB...")
all_data = collection.get(
# Ask for 750k (or more) to get everything
limit=750000,
include=["documents", "embeddings"]
)
documents = all_data['documents']
embeddings = all_data['embeddings']
from bertopic import BERTopic
print("Data fetched. Starting topic modeling...")
# We tell BERTopic to use our pre-made embeddings
topic_model = BERTopic(embedding_model="disable")
# This one-line command does it all:
# 1. Runs UMAP to reduce dimensions
# 2. Runs HDBSCAN to find clusters
# 3. Runs a c-TF-IDF to find topic words
topics, probs = topic_model.fit_transform(documents, embeddings)
print("Topic modeling complete!")
# 4. See your topics!
# This shows the most common topics
print(topic_model.get_topic_info())
# This shows the keywords for a specific topic
print(topic_model.get_topic(0))
topic_model.visualize_topics()