first-demo / get_bert_keywords.py
songhai2022's picture
Create get_bert_keywords.py
8902c6f
raw
history blame
488 Bytes
# import data manipulation libraries
from nltk.corpus import stopwords
from keybert import KeyBERT
from collections import Counter
# KeyBERT model
keybert_model = KeyBERT(model='all-MiniLM-L6-v2')
def get_bert_keywords(text, num_keywords, ignorewords):
result = keybert_model.extract_keywords(text, keyphrase_ngram_range=(
1, 1), top_n=num_keywords, stop_words=stopwords.words('english')+ignorewords)
keyberts = [keyword for keyword, score in result]
return keyberts