diff --git a/src/nlp.py b/src/nlp.py index 70c1e8d..c11c951 100644 --- a/src/nlp.py +++ b/src/nlp.py @@ -14,13 +14,10 @@ class NLP: tokens = [token.lemma_ for token in doc] return ' '.join(tokens) - def _get_nouns(self, tree: ParseTreeNode) -> List[str]: + def get_nouns(self, tree: ParseTree) -> List[str]: nouns: List[ParseTreeNode] = [] - for node in LevelOrderIter(tree): + for node in LevelOrderIter(tree.get_tree_root()): if node.upos != 'NOUN': continue nouns.append(node) - return [self._lemmatizer(noun.lemma) for noun in nouns] - - def get_nouns(self, tree: ParseTree) -> List[str]: - return self._get_nouns(tree.get_tree_root()) + return list(set([self.lemmatizer(noun.lemma) for noun in nouns]))