keyword_extraction.py

import nltk
import pytextrank
import re
from operator import itemgetter
import en_core_web_sm


class KeywordExtractor:
    """
    Keyword Extraction on text data
    Attributes:
        nlp: An instance English pipeline optimized for CPU for spacy
    """

    def __init__(self):
        self.nlp = en_core_web_sm.load()
        self.nlp.add_pipe("textrank")

    def get_keywords(self, text, max_keywords):
        """
        Extract keywords from text.
        Parameters:
            text (str): The user input string to extract keywords from
        Returns:
            kws (list): list of extracted keywords
        """

        doc = self.nlp(text)

        kws = [i.text for i in doc._.phrases[:max_keywords]]

        return kws

    def get_keyword_indices(self, kws, text):
        """
        Extract keywords from text.
        Parameters:
            kws (list): list of extracted keywords
            text (str): The user input string to extract keywords from
        Returns:
            keyword_indices (list): list of indices for keyword boundaries in text
        """

        keyword_indices = []
        for s in kws:
            indices = [[m.start(), m.end()] for m in re.finditer(re.escape(s), text)]
            keyword_indices.extend(indices)

        return keyword_indices

    def merge_overlapping_indices(self, keyword_indices):
        """
        Merge overlapping keyword indices.
        Parameters:
            keyword_indices (list): list of indices for keyword boundaries in text
        Returns:
            keyword_indices (list): list of indices for keyword boundaries in with overlapping combined
        """

        # Sort the array on the basis of start values of intervals.
        keyword_indices.sort()

        stack = []
        # insert first interval into stack
        stack.append(keyword_indices[0])
        for i in keyword_indices[1:]:
            # Check for overlapping interval,
            # if interval overlap
            if (stack[-1][0] <= i[0] <= stack[-1][-1]) or (stack[-1][-1] == i[0]-1):
                stack[-1][-1] = max(stack[-1][-1], i[-1])
            else:
                stack.append(i)
        return stack

    def merge_until_finished(self, keyword_indices):
        """
        Loop until no overlapping keyword indices left.
        Parameters:
            keyword_indices (list): list of indices for keyword boundaries in text
        Returns:
            keyword_indices (list): list of indices for keyword boundaries in with overlapping combined
        """

        len_indices = 0
        while True:
            # Merge overlapping indices
            merged = self.merge_overlapping_indices(keyword_indices)
            # Check to see if merging reduced number of annotation indices
            # If merging did not reduce list return final indicies
            if len_indices == len(merged):
                out_indices = sorted(merged, key=itemgetter(0))
                return out_indices
            else:
                len_indices = len(merged)

    def get_annotation(self, text, keyword_indices):
        """
        Create text annotation for extracted keywords.
        Parameters:
            keyword_indices (list): list of indices for keyword boundaries in text
        Returns:
            annotation (list): list of tuples for generating html
        """

        # Turn list to numpy array
        arr = list(text)

        # Loop through indices in list and insert delimeters
        for idx in sorted(keyword_indices, reverse=True):
            arr.insert(idx[0], "<kw>")
            arr.insert(idx[1]+1, "<!kw> <kw>")

        # join array
        joined_annotation = ''.join(arr)

        # split array on delimeter
        split = joined_annotation.split('<kw>')

        # Create annotation for keywords in text
        annotation = [(x.replace('<!kw> ', ''), "KEY", "#26aaef") if "<!kw>" in x else x for x in split]

        return annotation

    def generate(self, text, max_keywords):
        """
        Create text annotation for extracted keywords.
        Parameters:
            text (str): The user input string to extract keywords from
            max_keywords (int): Limit on number of keywords to generate
        Returns:
            annotation (list): list of tuples for generating html
            kws (list): list of extracted keywords
        """

        kws = self.get_keywords(text, max_keywords)

        indices = list(self.get_keyword_indices(kws, text))
        if indices:
            indices_merged = self.merge_until_finished(indices)
            annotation = self.get_annotation(text, indices_merged)
        else:
            annotation = None

        return annotation, kws