From 2e55579bb5e1f6ec40c85e79e41955a74336099f Mon Sep 17 00:00:00 2001
From: Aleksey Filippov <al.filippov@ulstu.ru>
Date: Thu, 23 Jun 2022 22:21:17 +0400
Subject: [PATCH] Add simple template support

---
 main.py    | 6 +++---
 src/nlp.py | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

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]))