deepset-ai / haystack

:mag: AI orchestration framework to build customizable, production-ready LLM applications. Connect components (models, vector DBs, file converters) to pipelines or agents that can interact with your data. With advanced retrieval methods, it's best suited for building RAG, question answering, semantic search or conversational agent chatbots.
https://haystack.deepset.ai
Apache License 2.0
17.18k stars 1.88k forks source link

Implement keyword extraction with KeyBERT #7083

Open JohnnyRacer opened 7 months ago

JohnnyRacer commented 7 months ago

Hello, I think adding a keyword extractor with KeyBERT would be quite useful. The keywords extracted could be used for paraphrasing or summarizing with logit_bias to allow for more consistent word usage in those tasks. Alternatively the keywords can also be used for locating a set documents from a collection that matches the keywords or phrases (sparse keyword based retrieval) . I have created snippet below based on the usage section KeyBERT's README.MD and the example from NamedEntityExtractor below.

from haystack.dataclasses import Document
from haystack.components.extractors import KeywordExtractor

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).
      """

extractor = KeywordExtractor(model="all-MiniLM-L6-v2") # Compatible with all 'sentence-transformer' models 

documents = [Document(content=doc)]

extractor.warm_up()
keywords_output = extractor.run(documents, keyphrase_ngram_range=(1, 1), stop_words=None)
# Longer phrases instead of keywords can be extracted by altering the keyphrase_ngram_range parameter
print(keywords_output)

"""
Expected output from the 'keywords_output':

[
    Document(id=97fc47fdd6aeb2540d0b015b234088b7386abea93671a7d336a80c244387457a,
    content: doc, # The input document from above
    meta: { # Adds the keywords to the document metadata
        'keywords' : [('learning', 0.4604),
        ('algorithm', 0.4556),
        ('training', 0.4487),
        ('class', 0.4086),
        ('mapping', 0.3700)]
    }
]

"""
masci commented 7 months ago

Hi @JohnnyRacer and thanks for the details you put in the issue!

We talked about this internally and while we understand the use case we couldn't figure out a way to prioritise this work. I imagine writing a custom component would be a good workaround in the meantime, but I'm also labelling this issue as contributions wanted in case anybody want to give it a shot.