POS Tag informed Lemmatizer

pos-lemmatizer.py

class POSTagLemmatizer:
    """Wrapper around NLTK's WordNetLemmatizer that takes into
    account token's POS tag."""

    ABBR_TO_TAG = {
        "n": ["NN", "NNS", "NNP", "NNPS"],
        "v": ["VB", "VBD", "VBG", "VBN", "VBP", "VBZ"],
        "r": ["RB", "RBR", "RBS"],
        "a": ["JJ", "JJR", "JJS"],
    }

    TAG_TO_ABBR = {tag: abb for abb, tags in ABBR_TO_TAG.items() for tag in tags}

    def __init__(self):
        from nltk import pos_tag
        from nltk import word_tokenize
        from nltk.stem import WordNetLemmatizer

        self._tokenize = word_tokenize
        self._lemmatizer = WordNetLemmatizer()
        self._get_pos_tags = pos_tag

    def lemmatize(self, text):
        tokens = self._tokenize(text)
        pos_tags = self._get_pos_tags(tokens)

        return [
            self._lemmatizer.lemmatize(tkn, self.get_pos_tag_lw(tag))
            for tkn, tag in pos_tags
        ]

    @classmethod
    def get_pos_tag_lw(cls, pos_tag):
        return cls.TAG_TO_ABBR.get(pos_tag, "n")

    def __call__(self, tokens):
        return self.lemmatize(tokens)