Merge tag '0.8.2' into develop

no message
ckiplab · Apr 28, 2020 · b0b17d1 · b0b17d1
2 parents e034028 + d440427
commit b0b17d1
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/ckipnlp/__init__.py b/ckipnlp/__init__.py
@@ -10,7 +10,7 @@
 __copyright__ = '2018-2020 CKIP Lab'
 
 __title__ = 'CKIPNLP'
-__version__ = '0.8.1'
+__version__ = '0.8.2'
 __description__ = 'CKIP CoreNLP'
 __license__ = 'CC BY-NC-SA 4.0'
 

diff --git a/ckipnlp/driver/coref.py b/ckipnlp/driver/coref.py
@@ -158,11 +158,11 @@ def _get_result(cls, tree_list, *, coref_tree):
 
         # Assign coref ID
         node2coref = {}  # (tree_id, node_id) => ref_id
-        coref2node = {}  # ref_id => node
+        sources = set()  # (tree_id, node_id)
 
         for ref_id, coref_source in enumerate(coref_tree.children(coref_tree.root)):
             tree_id, node_id = coref_source.identifier
-            coref2node[ref_id] = tree_list[tree_id][node_id]
+            sources.add((tree_id, node_id,))
             for tree_id, node_id in coref_tree.expand_tree(coref_source.identifier):
                 node2coref[tree_id, node_id] = ref_id
 
@@ -201,11 +201,10 @@ def _get_result(cls, tree_list, *, coref_tree):
             for node in nodes:
                 ref_id = node2coref.get((tree_id, node.identifier,), -1)
                 if ref_id >= 0:
-                    ref_node = coref2node[ref_id]
                     tokens.append(_CorefToken(  # pylint: disable=no-value-for-parameter
                         word=node.data.word,
                         idx=node.identifier,
-                        coref=(ref_id, 'source' if node.identifier == ref_node.identifier else 'target',),
+                        coref=(ref_id, 'source' if (tree_id, node.identifier,) in sources else 'target',),
                     ))
                 else:
                     tokens.append(_CorefToken(  # pylint: disable=no-value-for-parameter