diff --git a/combo/utils/graph.py b/combo/utils/graph.py index ed2c1ff1c81cac086a59942168d153d1a23d24a7..b4c9632811d34037d1155d0c349480e2276cc738 100644 --- a/combo/utils/graph.py +++ b/combo/utils/graph.py @@ -114,13 +114,31 @@ def restore_collapse_edges(tree_tokens): head, relation = d.split(':', 1) ehead = f"{len(tree_tokens)}.{len(empty_tokens) + 1}" empty_node_relation, current_node_relation = relation.split(">", 1) - deps[i] = f"{ehead}:{current_node_relation}" - empty_tokens.append( - { - "id": ehead, - "deps": f"{head}:{empty_node_relation}" - } - ) + # Edge case, double > + if ">" in current_node_relation: + second_empty_node_relation, current_node_relation = current_node_relation.split(">") + deps[i] = f"{ehead}:{current_node_relation}" + empty_tokens.append( + { + "id": ehead, + "deps": f"{head}:{empty_node_relation}" + } + ) + empty_tokens.append( + { + "id": f"{len(tree_tokens)}.{len(empty_tokens) + 1}", + "deps": f"{ehead}:{second_empty_node_relation}" + } + ) + + else: + deps[i] = f"{ehead}:{current_node_relation}" + empty_tokens.append( + { + "id": ehead, + "deps": f"{head}:{empty_node_relation}" + } + ) deps = sorted([d.split(":", 1) for d in deps], key=lambda x: float(x[0])) token["deps"] = "|".join([f"{k}:{v}" for k, v in deps]) return empty_tokens