Suggestion & Feedback: Add support for Apple Silicons' GPU support

[junxi-liu]

First of all thanks a lot for this amazing package! I am currently using a MacBook Pro 2021 version with M1 Pro chip. As many Apple users may know, PyTorch now supports GPU acceleration by using mps as computing device. However, currently in the BERTopic model, there is no support for this parameter.

After digging into the details for some time I found a easy work around: SentenceTransformer do supports the device parameter, so we can simply do the following:

embedding = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2', device = 'mps')

model = BERTopic(language="chinese (simplified)", embedding_model=embedding, nr_topics='auto',verbose=True, vectorizer_model=vectorizer,)

which will perform exactly as the default model if an user chooses to. On my Mac, this drastically reduce the training time from 4h30m to 58m.

I assume this should be also easy to incorporate into the package, which will be very helpful. Thanks!

[MaartenGr]

Thanks for the suggestion! It is indeed troublesome that the right device is nog chosen by default it seems. I am not sure though how that would be integrated within BERTopic since it is relatively minor to give it a parameter to tweak. For example, it would only work for SentenceTransformers and not necessarily for all other embedding models. I want to limit the parameter space as much as possible to prevent an overcomplicated user experience. Perhaps a mention in the documentation would be more appropriate?

[netique]

I realized that sentence_transformers utilize pytorch that supports mps by accident. A documentation entry may suffice indeed. But maybe this all concerns sentence_transformers not bertopic necessarily.

[MaartenGr]

There is a PR open that is relevant to this discussion: https://github.com/UKPLab/sentence-transformers/pull/2342 When this PR is merged, it should automatically select the right device.

[tomaarsen]

I've merged https://github.com/UKPLab/sentence-transformers/pull/2342. I intend to release within the next weeks to include it into SentenceTransformers 2.3.0. Hope that helps.

On an unrelated note, much respect for you and your work @MaartenGr.

• Tom Aarsen

[MaartenGr]

@tomaarsen Thanks for replying here! It will definitely help users wanting an easier workflow. Thank you for the kind words! Glad to have you on the sentence-transformers package. You're an amazing engineer and I'm excited to see your work in the sentence-transformers package.