MaartenGr / KeyBERT

Minimal keyword extraction with BERT
https://MaartenGr.github.io/KeyBERT/
MIT License
3.46k stars 344 forks source link

KeyLLM fails when no GPU is available #196

Closed linards-kalvans closed 8 months ago

linards-kalvans commented 9 months ago

KeyLLM fails when no GPU is available: `

RuntimeError Traceback (most recent call last) Cell In[14], line 13 10 kw_model = KeyBERT(llm = llm) 12 # Extract keywords ---> 13 keywords = kw_model.extract_keywords(sparks.query('id == "4d430272-7668-4353-b113-d1ebdda308a4"')['full_text'].values[0])

File ~/Dokumenti/Python_projects/DS-prototypes/Spark summarization/.venv/lib/python3.11/site-packages/keybert/_model.py:259, in KeyBERT.extract_keywords(self, docs, candidates, keyphrase_ngram_range, stop_words, top_n, min_df, use_maxsum, use_mmr, diversity, nr_candidates, vectorizer, highlight, seed_keywords, doc_embeddings, word_embeddings, threshold) 257 if self.llm is not None: 258 import torch --> 259 doc_embeddings = torch.from_numpy(doc_embeddings).float().to("cuda") 260 if isinstance(all_keywords[0], tuple): 261 candidate_keywords = [[keyword[0] for keyword in all_keywords]]

File ~/Dokumenti/Python_projects/DS-prototypes/Spark summarization/.venv/lib/python3.11/site-packages/torch/cuda/init.py:298, in _lazy_init() 296 if "CUDA_MODULE_LOADING" not in os.environ: 297 os.environ["CUDA_MODULE_LOADING"] = "LAZY" --> 298 torch._C._cuda_init() 299 # Some of the queued calls may reentrantly call _lazy_init(); 300 # we need to just return without initializing in that case. 301 # However, we must not let any other threads in! 302 _tls.is_initializing = True

RuntimeError: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx `

Seems that the issue is here: https://github.com/MaartenGr/KeyBERT/blob/6638d44e7c133b43f092a6480696dc14491bb279/keybert/_model.py#L259 If I remove .to("cuda") everything works like a charm.

MaartenGr commented 9 months ago

Hmmm, not quite sure what the best solution would be here. Perhaps an automatic detection of the system and using that instead. I imagine mac users would run into similar issues.