Extract entities from sentences

Question

I have a function that extracts basic parts-of-speech words (NN,NNP,JJ,VBS etc.) from about 2000 sentences. I want to understand if there is a way to optimize so I can bring down the execution time (from about 12 secs right now to single digits) if possible.

import re, nltk, html

def extract_entities(text_list):

    entities=set()

    for sentence in text_list:
        try:

            tokens = nltk.word_tokenize(sentence)
            tagged = nltk.pos_tag(tokens)
            [entities.add(word.lower()) for word,pos in tagged
                 if pos in ['NN','NNP','NNS','NNPS','JJ','VBS','VBN','VBG','VBD']]

        except Exception as e:
            print('stage 1',e)

    return entities

Answer 1

Here are a few possible improvements:

1. Use a set for the positions you want and make it a constant. (This will give you a small speed-boost, because in is O(n) for lists, but O(1) for sets.)

2. Don't use a list-comprehension with side-effects for a loop.

3. Follow Python's official style-guide, PEP8. Specifically:

   a) Don't do multiple imports

   b) Put spaces around your operators

   c) Put a space after a comma in an argument list

4. Limit the scope of your try..except

With these fixes, here is the final code, as I would write it:

import re
import nltk
import html

POSITIONS = {'NN','NNP','NNS','NNPS','JJ','VBS','VBN','VBG','VBD'}

def extract_entities(text_list):
    entities = set()
    for sentence in text_list:
        try:
            tokens = nltk.word_tokenize(sentence)
            tagged = nltk.pos_tag(tokens)
        except Exception as e:
            print('stage 1', e)
            continue
        for word, pos in tagged:
            if pos in POSITIONS:
                entities.add(word.lower())
    return entities