KeyBERT
KeyBERT copied to clipboard
For tens of millions of data, how to use GPU to accelerate the KeyBERT algorithm?
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.
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