KeyBERT Documentation

Use the KeyBERT.extract_keywords() method to extract keywords from a document.

Key parameters:

• keyphrase_ngram_range: A tuple defining the (min, max) length of N-grams. Use (1, 1) for single words and (1, 2) or higher for keyphrases.
• stop_words: Specify stop words (e.g., 'english') to filter results.
• highlight: Set to True to highlight the extracted keywords in the original document text.

Recommended models:

• English: all-MiniLM-L6-v2
• Multilingual: paraphrase-multilingual-MiniLM-L12-v2

from keybert import KeyBERT

doc = """
         Supervised learning is the machine learning task of learning a function that
         maps an input to an output based on example input-output pairs. It infers a
         function from labeled training data consisting of a set of training examples.
         In supervised learning, each example is a pair consisting of an input object
         (typically a vector) and a desired output value (also called the supervisory signal).
         A supervised learning algorithm analyzes the training data and produces an inferred function,
         which can be used for mapping new examples. An optimal scenario will allow for the
         algorithm to correctly determine the class labels for unseen instances. This requires
         the learning algorithm to generalize from the training data to unseen situations in a
         'reasonable' way (see inductive bias).
      ""
kw_model = KeyBERT()

# Extract single words
keywords = kw_model.extract_keywords(doc, keyphrase_ngram_range=(1, 1), stop_words=None)

# Extract keyphrases
keyphrases = kw_model.extract_keywords(doc, keyphrase_ngram_range=(1, 2), stop_words=None)

# Highlight keywords
highlighted_doc = kw_model.extract_keywords(doc, highlight=True)

KeyBERT uses CountVectorizer to split documents into candidate keywords or keyphrases. Since splitting occurs after embedding, you can customize the vectorizer to change how candidates are parsed without affecting embedding quality. You can pass a CountVectorizer instance to the extract_keywords method.

from keybert import KeyBERT
from sklearn.feature_extraction.text import CountVectorizer

doc = "Your document text here"
kw_model = KeyBERT()
vectorizer = CountVectorizer()
keywords = kw_model.extract_keywords(doc, vectorizer=vectorizer)

from keybert import KeyBERT
from sklearn.feature_extraction.text import CountVectorizer

doc = """
         Supervised learning is the machine learning task of learning a function that
         maps an input to an output based on example input-output pairs.[1] It infers a
         function from labeled training data consisting of a set of training examples.[2]
         In supervised learning, each example is a pair consisting of an input object
         (typically a vector) and a desired output value (also called the supervisory signal).
         A supervised learning algorithm analyzes the training data and produces an inferred function,
         which can be used for mapping new examples. An optimal scenario will allow for the
         algorithm to correctly determine the class labels for unseen instances. This requires
         the learning algorithm to generalize from the training data to unseen situations in a
         'reasonable' way (see inductive bias).
      ""
kw_model = KeyBERT()
vectorizer = CountVectorizer()
keywords = kw_model.extract_keywords(doc, vectorizer=vectorizer)

Spacy models can be used as embedding backends. It is recommended to exclude unnecessary components like tagger, parser, etc., to save resources.

To use a standard Spacy model:

import spacy

nlp = spacy.load("en_core_web_md", exclude=['tagger', 'parser', 'ner', 'attribute_ruler', 'lemmatizer'])
kw_model = KeyBERT(model=nlp)

To use a Spacy transformer model (_trf):

import spacy

spacy.prefer_gpu()
nlp = spacy.load("en_core_web_trf", exclude=['tagger', 'parser', 'ner', 'attribute_ruler', 'lemmatizer'])
kw_model = KeyBERT(model=nlp)

If you encounter memory issues with transformer models, configure the GPU allocator:

import spacy
from thinc.api import set_gpu_allocator, require_gpu

nlp = spacy.load("en_core_web_trf", exclude=['tagger', 'parser', 'ner', 'attribute_ruler', 'lemmatizer'])
set_gpu_allocator("pytorch")
require_gpu(0)

kw_model = KeyBERT(model=nlp)

Use KeyLLM to perform keyword extraction via LLMs (e.g., OpenAI).

Setup:

1. Install the OpenAI package: pip install openai.
2. Create an OpenAI client.
3. Wrap the client in keybert.llm.OpenAI.
4. Initialize KeyLLM with the LLM object.

Efficient Extraction: To avoid redundant LLM calls, you can provide document embeddings and a threshold. KeyLLM will only query the LLM for documents that are sufficiently different from previously processed ones.

Parameters for extract_keywords in KeyLLM:

• embeddings: Pre-computed embeddings for the documents.
• threshold: Similarity threshold to decide if documents should share the same keywords.

import openai
from keybert.llm import OpenAI
from keybert import KeyLLM
from sentence_transformers import SentenceTransformer

# 1. Basic LLM Extraction
client = openai.OpenAI(api_key=MY_API_KEY)
llm = OpenAI(client)
kw_model = KeyLLM(llm)
keywords = kw_model.extract_keywords(doc)

# 2. Efficient Extraction with Embeddings
model = SentenceTransformer('all-MiniLM-L6-v2')
embeddings = model.encode(MY_DOCUMENTS, convert_to_tensor=True)

kw_model = KeyLLM(llm)
keywords = kw_model.extract_keywords(MY_DOCUMENTS, embeddings=embeddings, threshold=.75)