diff --git a/main.py b/main.py index c40604a..ad33283 100644 --- a/main.py +++ b/main.py @@ -15,9 +15,9 @@ def _main(wav_file: str): print(f'Text: {text}') parse_tree: ParseTree = Syntax().get_parse_tree(text) print(f'Parse tree:\n{parse_tree}') - nouns: List[str] = NLP().get_nouns(parse_tree) - print(f'Nouns:\n{" ".join(nouns)}') - result: str = MyOntology().get_event_description(nouns) + terms: List[str] = NLP().get_by_template(parse_tree, ['NOUN', 'VERB', 'ADJ']) + print(f'Extracted terms:\n{" ".join(terms)}') + result: str = MyOntology().get_event_description(terms) print(f'Test:\n{result}') diff --git a/src/nlp.py b/src/nlp.py index 1b55aa7..b311347 100644 --- a/src/nlp.py +++ b/src/nlp.py @@ -13,10 +13,10 @@ class NLP: tokens = [token.lemma_ for token in doc] return ' '.join(tokens) - def get_nouns(self, tree: ParseTree) -> List[str]: + def get_by_template(self, tree: ParseTree, template: List[str]) -> List[str]: nouns: List[ParseTreeNode] = [] for node in LevelOrderIter(tree.get_tree_root()): - if node.upos != 'NOUN': + if node.upos not in template: continue nouns.append(node) return list(set([self.lemmatizer(noun.lemma) for noun in nouns]))