KeyBERT icon indicating copy to clipboard operation
KeyBERT copied to clipboard

For tens of millions of data, how to use GPU to accelerate the KeyBERT algorithm?

Open tmacsyf opened this issue 9 months ago • 1 comments

Hello author, since I have millions of documents, I want to use GPU to accelerate the computation of the KeyBERT algorithm, especially during the embedding phase. However, I don’t see a place to set the batch_size parameter. Are there any methods to speed up the computation? Thank you.

tmacsyf avatar May 09 '25 03:05 tmacsyf

GPU acceleration is enabled by default when using sentence-transformers. There's currently no access to the batch_size parameter in sentence-transformers, but you can change it here: https://github.com/MaartenGr/KeyBERT/blob/master/keybert/backend/_sentencetransformers.py

MaartenGr avatar May 12 '25 10:05 MaartenGr