MaartenGr / MaartenGr/BERTopic
Saving a trained model using pytorch and safetensor and then redownloading causes topics to be off
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 7.8k
- Forks
- 920
- Avg merge
- 22h 24m
- Merged PRs (30d)
- 5
Description
### Have you searched existing issues? 🔎
- [x] I have searched and found no existing issues
### Desribe the bug
After training, I tried saving the model using both pytorch and safetensor. When I re-download the model, load the files into Bertopic using Bertopic.load(), and run inference using transform(), all the topics are coming out differently than the original fit results. Below are some examples the first topic and prob is from the original training/fit of the model and the second is from running transform():
Topic: 2 Probability: 0.9999999985560923 vs. Topic: 3 Probability: 0.9999477863311768
Topic: 1 Probability: 0.9993163446248252 vs. Topic: 2 Probability: 0.04614641437377926
Topic: 2 Probability: 1.0 vs. Topic: 3 Probability: 0.9591490626335144
One thing to note is that running transform over and over comes out with the same results that are different than the original training output. Also, when I run transform on the original model without saving it anywhere else, I get the same results as the original run. I was wondering if I am missing something with saving the model correctly. Below is the code I use to train, save, and run transform on the model. We also run reduce_outliers() before saving the model.
### Reproduction
```python
from bertopic import BERTopic
self.model_params = {
'min_topic_size': int((len(rows) / 160) - 1),
'calculate_probabilities': True,
'verbose': True,
'umap_model': UMAP(
n_neighbors=50,
n_components=20,
metric='cosine',
low_memory=False,
random_state=42,
)
}
self.model = BERTopic(**self.model_params)
self.topics, self.probabilities = self.model.fit_transform(
documents=self.docs,
embeddings=numpy.array(self.embeddings),
y=self.labels
)
new_topics = self.model.reduce_outliers(
self.docs,
self.topics,
probabilities=self.probabilities,
strategy='probabilities'
)
self.model.update_topics(self.docs, topics=new_topics)
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
self.model.save(torch_file_path, serialization="safetensors", save_ctfidf=True, save_embedding_model=embedding_model)
new_model = BERTopic.load(artifact_path)
new_model_temp_topics, new_model_temp_probabilities = new_model.transform(documents=self.docs, embeddings=numpy.array(self.embeddings))
```
### BERTopic Version
0.16.0
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Reproduce the reported flow from BERTopic.fit_transform through reduce_outliers(), update_topics(), save(), BERTopic.load(), and transform() using the supplied code and version 0.16.0. Compare the original and reloaded model topics and probabilities, including both PyTorch and safetensors serialization. Done means reloaded-model inference matches the original results or the discrepancy is documented with a focused regression test.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100