Release 1.0.4

combo/predict.py

@@ -197,14 +197,15 @@ class COMBO(predictor.Predictor):
 
         if "enhanced_head" in predictions and predictions["enhanced_head"]:
             # TODO off-by-one hotfix, refactor
-            h = np.array(predictions["enhanced_head"])
+            sentence_length = len(tree_tokens)
+            h = np.array(predictions["enhanced_head"])[:sentence_length, :sentence_length]
             h = np.concatenate((h[-1:], h[:-1]))
-            r = np.array(predictions["enhanced_deprel_prob"])
+            r = np.array(predictions["enhanced_deprel_prob"])[:sentence_length, :sentence_length, :]
             r = np.concatenate((r[-1:], r[:-1]))
 
             graph.graph_and_tree_merge(
-                tree_arc_scores=predictions["head"],
-                tree_rel_scores=predictions["deprel"],
+                tree_arc_scores=predictions["head"][:sentence_length],
+                tree_rel_scores=predictions["deprel"][:sentence_length],
                 graph_arc_scores=h,
                 graph_rel_scores=r,
                 idx2label=self.vocab.get_index_to_token_vocabulary("deprel_labels"),