Refine node tree define and print.

wit/dataset/meaning_dataset.py

@@ -190,7 +190,8 @@ class MeaningMap:
                 tree[m] = get_tree_dict(ms_map, vocab_size, m) if m >= self.vocab_size else m
             return tree
 
-        return get_tree_dict(self.ms_map, self.vocab_size, meaning)
+        td = get_tree_dict(self.ms_map, self.vocab_size, meaning)
+        return {meaning: td}
 
     def max_length(self):
         return max(self.ms_len)

wit/dataset/node_tree.py

@@ -1,4 +1,21 @@
-from anytree import Node, RenderTree
+from anytree import RenderTree
+
+from anytree.node.nodemixin import NodeMixin
+from anytree.node.util import _repr
+
+
+class Node(NodeMixin):
+    def __init__(self, name, parent=None, children=None, **kwargs):
+        self.__dict__.update(kwargs)
+        self.name = name
+        self.prop = ""
+        self.parent = parent
+        if children:
+            self.children = children
+
+    def __repr__(self):
+        args = ["%r" % self.separator.join([""] + [str(node.name) for node in self.path])]
+        return _repr(self, args=args, nameblacklist=["name"])
 
 
 class NodeTree:
@@ -6,6 +23,7 @@ class NodeTree:
         self.tree = tree
         self.node = None
         self.seq_node = []
+        self.get_node()
 
     def get_node(self):
         def get_tree_node(tree, parent, seqlist):
@@ -25,10 +43,13 @@ class NodeTree:
         self.node = root.children[0] if len(self.tree) == 1 else root
         return self.node
 
+    def set_seq_prop(self, index, prop):
+        self.seq_node[index].prop = prop
+
     def print(self):
         treestr = ""
         for pre, fill, node in RenderTree(self.get_node()):
-            treestr += f"{pre}{node.name}\n"
+            treestr += f"{pre}{node.name} {node.prop}\n"
         print(treestr)
 
     def print_tree(tree):