using external vocabulary for synonyms and antonyms

kocmitom · kocmitom · commit a524c3fad393 · 2017-12-10T16:46:27.000+01:00
diff --git a/udapi/block/newspeak/prevele.py b/udapi/block/newspeak/prevele.py
@@ -21,13 +21,23 @@
 class PreVele(Block):
     """Change all comparatives to vele-x and superlatives to převele-x."""
 
-    def __init__(self, morphodita_path='models/morphodita/cs/',
+    def __init__(self, synonyms, antonyms, morphodita_path='models/morphodita/cs/',
                  morphodita_model='czech-morfflex-131112.dict',
                  **kwargs):
         """Create the PreVele block object."""
         super().__init__(**kwargs)
+        self.antonyms = self.load_vocab(antonyms)
+        self.synonyms = self.load_vocab(synonyms)
         self.morphodita = MorphoDiTa(model=morphodita_path + morphodita_model)
 
+    def load_vocab(self, filename):
+        vocab = {}
+        with open(filename) as f:
+            for line in f:
+                (key, val) = line.strip().split('\t')
+                vocab[key] = val
+        return vocab
+
     def process_tree(self, tree):
 
         # apply process_node on all nodes
@@ -42,7 +52,7 @@ def process_tree(self, tree):
         tree.text = tree.compute_text()
 
     def process_node(self, node):
-        antonym = ANTONYMS.get(node.lemma)
+        antonym = self.antonyms.get(node.lemma)
         if antonym is not None:
             if node.xpos[11] == 'N':
                 if node.form.lower().startswith('ne'):
@@ -56,6 +66,14 @@ def process_node(self, node):
                     node.xpos = node.xpos[:10] + 'N' + node.xpos[11:]
                     node.form = 'ne' + forms[0].form
 
+        synonym = self.synonyms.get(node.lemma)
+        if synonym is not None:
+            forms = self.morphodita.forms_of_lemma(synonym, node.xpos)
+            if forms:
+                node.lemma = synonym
+                node.form = forms[0].form
+  
+
         degree = node.feats["Degree"]
         if degree in ("Sup", "Cmp"):
             new_xpos = node.xpos[:9] + '1' + node.xpos[10:]
