Skip to content

Serialization

Saving, loading, and sharing a BERTopic model can be done in several ways. It is generally advised to go with .safetensors as that allows for a small, safe, and fast method for saving your BERTopic model. However, other formats, such as .pickle and pytorch .bin are also possible.

Saving

There are three methods for saving BERTopic:

  1. A light model with .safetensors and config files
  2. A light model with pytorch .bin and config files
  3. A full model with .pickle

Tip

It is advised to use methods 1 or 2 for saving as they generated very small models. Especially method 1 (safetensors) allows for a relatively safe format compared to the other methods.

The methods are used as follows:

topic_model = BERTopic().fit(my_docs)

# Method 1 - safetensors
embedding_model = "sentence-transformers/all-MiniLM-L6-v2"
topic_model.save("path/to/my/model_dir", serialization="safetensors", save_ctfidf=True, save_embedding_model=embedding_model)

# Method 2 - pytorch
embedding_model = "sentence-transformers/all-MiniLM-L6-v2"
topic_model.save("path/to/my/model_dir", serialization="pytorch", save_ctfidf=True, save_embedding_model=embedding_model)

# Method 3 - pickle
topic_model.save("my_model", serialization="pickle")

Warning

When saving the model, make sure to also keep track of the versions of dependencies and Python used. Loading and saving the model should be done using the same dependencies and Python. Moreover, models saved in one version of BERTopic are not guaranteed to load in other versions.

Pickle Drawbacks

Saving the model with pickle allows for saving the entire topic model, including dimensionality reduction and clustering algorithms, but has several drawbacks:

  • Arbitrary code can be run from .pickle files
  • The resulting model is rather large (often > 500MB) since all sub-models need to be saved
  • Explicit and specific version control is needed as they typically only run if the environment is exactly the same

Safetensors and Pytorch Advantages

Saving the topic modeling with .safetensors or pytorch has a number of advantages:

  • .safetensors is a relatively safe format
  • The resulting model can be very small (often < 20MB>) since no sub-models need to be saved
  • Although version control is important, there is a bit more flexibility with respect to specific versions of packages
  • More easily used in production
  • Share models with the HuggingFace Hub





The above image, a model trained on 100,000 documents, demonstrates the differences in sizes comparing safetensors, pytorch, and pickle. The difference in sizes can mostly be explained due to the efficient saving procedure and that the clustering and dimensionality reductions are not saved in safetensors/pytorch since inference can be done based on the topic embeddings.

HuggingFace Hub

When you have created a BERTopic model, you can easily share it with other through the HuggingFace Hub. First, you need to log in to your HuggingFace account which you can do in a number of ways:

  • Log in to your Hugging Face account with the command below
huggingface-cli login

# or using an environment variable
huggingface-cli login --token $HUGGINGFACE_TOKEN
  • Alternatively, you can programmatically login using login() in a notebook or a script
from huggingface_hub import login
login()
  • Or you can give a token with the token variable

When you have logged in to your HuggingFace account, you can save and upload the model as follows:

from bertopic import BERTopic

# Train model
topic_model = BERTopic().fit(my_docs)

# Push to HuggingFace Hub
topic_model.push_to_hf_hub(
    repo_id="MaartenGr/BERTopic_ArXiv",
    save_ctfidf=True
)

# Load from HuggingFace
loaded_model = BERTopic.load("MaartenGr/BERTopic_ArXiv")

Parameters

There are number of parameters that may be worthwhile to know:

  • private
    • Whether to create a private repository
  • serialization
    • The type of serialization. Either safetensors or pytorch. Make sure to run pip install safetensors for safetensors.
  • save_embedding_model
    • A pointer towards a HuggingFace model to be loaded in with SentenceTransformers. E.g., sentence-transformers/all-MiniLM-L6-v2
  • save_ctfidf
    • Whether to save c-TF-IDF information

Loading

To load a model:

# Load from directory
loaded_model = BERTopic.load("path/to/my/model_dir")

# Load from file
loaded_model = BERTopic.load("my_model")

# Load from HuggingFace
loaded_model = BERTopic.load("MaartenGr/BERTopic_Wikipedia")

The embedding model cannot always be saved using a non-pickle method if, for example, you are using OpenAI embeddings. Instead, you can load them in as follows:

# Define embedding model
import openai
from bertopic.backend import OpenAIBackend

client = openai.OpenAI(api_key="sk-...")
embedding_model = OpenAIBackend(client, "text-embedding-ada-002")

# Load model and add embedding model
loaded_model = BERTopic.load("path/to/my/model_dir", embedding_model=embedding_model)