feat(#87): switch to spacy for lemmatization

Author: jannis-baum

script/README.md

@@ -1,4 +1,5 @@
 # Python interface to TEM
 
 - install `requirements.txt`
-- `python3 -m nltk.download punkt stopwords`
+- `python3 -m nltk.download punkt`
+- for English: `python -m spacy download en_core_web_md`

script/nlp.py

@@ -1,8 +1,7 @@
 import re
 
-from nltk.corpus import wordnet
-from nltk.tokenize import word_tokenize, sent_tokenize
-from nltk import WordNetLemmatizer, pos_tag
+from nltk.tokenize import sent_tokenize
+import spacy
 
 # ------------------------------------------------------------------------------
 # MARK: normalization
@@ -28,26 +27,9 @@ def normalize_quotes(doc: str) -> str:
 def normal_str(doc: str) -> str:
     return lower_alnum(normalize_quotes(doc))
 
-__lemma = WordNetLemmatizer()
-
+__lemma = spacy.load('en_core_web_md')
 def normal_tokens(doc: str) -> list[str]:
-    def __get_wordnet_pos(treebank_tag):
-        if treebank_tag.startswith('J'):
-            return wordnet.ADJ
-        elif treebank_tag.startswith('V'):
-            return wordnet.VERB
-        elif treebank_tag.startswith('N'):
-            return wordnet.NOUN
-        elif treebank_tag.startswith('R'):
-            return wordnet.ADV
-        else:
-            return wordnet.NOUN
-
-    token_pos = pos_tag(word_tokenize(normal_str(doc)))
-    return [
-        __lemma.lemmatize(token, __get_wordnet_pos(pos))
-        for token, pos in token_pos
-    ]
+    return [token.lemma_ for token in __lemma(normal_str(doc))]
 
 # ------------------------------------------------------------------------------
 # MARK: TEM prep

script/requirements.txt

@@ -3,3 +3,6 @@ nltk>=3.6.7
 numpy>=1.23.5
 pexpect>=4.8.0
 pyyaml>=6.0
+spacy>=3.7.4
+spacy-legacy>=3.0.12
+spacy-loggers>=1.0.5