from bertopic import BERTopic
from sklearn.datasets import fetch_20newsgroups
 
docs = fetch_20newsgroups(subset='all',  remove=('headers', 'footers', 'quotes'))['data']

"""
Explanation: BERTopic - Tutorial
We start with installing bertopic from pypi before preparing the data. 
NOTE: Make sure to select a GPU runtime. Otherwise, the model can take quite some time to create the document embeddings!
Prepare data
For this example, we use the popular 20 Newsgroups dataset which contains roughly 18000 newsgroups posts on 20 topics.
End of explanation
"""


model = BERTopic(language="english")
topics, probs = model.fit_transform(docs)

"""
Explanation: Create Topics
We select the "english" as the main language for our documents. If you want a multilingual model that supports 50+ languages, please select "multilingual" instead.
End of explanation
"""


model.get_topic_freq().head(5)

"""
Explanation: We can then extract most frequent topics:
End of explanation
"""


model.get_topic(49)[:10]

"""
Explanation: -1 refers to all outliers and should typically be ignored. Next, let's take a look at the most frequent topic that was generated:
End of explanation
"""


from bertopic import languages
print(languages)

"""
Explanation: Note that the model is stocastich which mmeans that the topics might differ across runs. 
For a full list of support languages, see the values below:
End of explanation
"""


# st_model = BERTopic(embedding_model="xlm-r-bert-base-nli-stsb-mean-tokens")

"""
Explanation: Embedding model
You can select any model from sentence-transformers and use it instead of the preselected models by simply passing the model through
BERTopic with embedding_model:
End of explanation
"""


model.visualize_topics()

"""
Explanation: Click here for a list of supported sentence transformers models.  
Visualize Topics
After having trained our BERTopic model, we can iteratively go through perhaps a hundred topic to get a good 
understanding of the topics that were extract. However, that takes quite some time and lacks a global representation. 
Instead, we can visualize the topics that were generated in a way very similar to 
LDAvis:
End of explanation
"""


model.visualize_distribution(probs[0])

"""
Explanation: Visualize Topic Probabilities
The variable probabilities that is returned from transform() or fit_transform() can 
be used to understand how confident BERTopic is that certain topics can be found in a document. 
To visualize the distributions, we simply call:
End of explanation
"""


new_topics, new_probs = model.reduce_topics(docs, topics, probs, nr_topics=60)

"""
Explanation: Topic Reduction
Finally, we can also reduce the number of topics after having trained a BERTopic model. The advantage of doing so, 
is that you can decide the number of topics after knowing how many are actually created. It is difficult to 
predict before training your model how many topics that are in your documents and how many will be extracted. 
Instead, we can decide afterwards how many topics seems realistic:
End of explanation
"""


model.update_topics(docs, topics, n_gram_range=(1, 3), stop_words="english")

"""
Explanation: The reasoning for putting docs, topics, and probs as parameters is that these values are not saved within 
BERTopic on purpose. If you were to have a million documents, it seems very inefficient to save those in BERTopic 
instead of a dedicated database.  
Topic Representation
When you have trained a model and viewed the topics and the words that represent them,
you might not be satisfied with the representation. Perhaps you forgot to remove
stop_words or you want to try out a different n_gram_range. We can use the function update_topics to update 
the topic representation with new parameters for c-TF-IDF:
End of explanation
"""


similar_topics, similarity = model.find_topics("vehicle", top_n=5); similar_topics

model.get_topic(28)

"""
Explanation: Search Topics
After having trained our model, we can use find_topics to search for topics that are similar 
to an input search_term. Here, we are going to be searching for topics that closely relate the 
search term "vehicle". Then, we extract the most similar topic and check the results:
End of explanation
"""


# Save model
model.save("my_model")

# Load model
my_model = BERTopic.load("my_model")

"""
Explanation: Model serialization
The model and its internal settings can easily be saved. Note that the documents and embeddings will not be saved. However, UMAP and HDBSCAN will be saved.
End of explanation
"""


