Upload BertForSyntaxParsing.py

BertForSyntaxParsing.py

@@ -189,7 +189,11 @@ class BertForSyntaxParsing(BertPreTrainedModel):
                 dep_func = ALL_FUNCTION_LABELS[funcs[i + 1]]
 
                 if dep_head == 'root': root_idx = len(tree)
-                tree.append(dict(word=toks[i], dep_head_idx=idx_mapping[dep_idx], dep_head=dep_head, dep_func=dep_func))
+                tree.append(dict(word=toks[i], dep_head_idx=idx_mapping[dep_idx], dep_func=dep_func))
+            # append the head word
+            for d in tree:
+                d['dep_head'] = tree[d['dep_head_idx']]['word']
+            
             outputs.append(dict(tree=tree, root_idx=root_idx))
         return outputs