From 038c97056cdbe50c6b60a24cb272d9826472ccb0 Mon Sep 17 00:00:00 2001
From: Mateusz Klimaszewski <mk.klimaszewski@gmail.com>
Date: Wed, 12 May 2021 17:01:30 +0200
Subject: [PATCH] Fix batch predictions for DEPS.
---
combo/predict.py | 9 +++++----
scripts/predict_iwpt21.py | 3 +++
2 files changed, 8 insertions(+), 4 deletions(-)
diff --git a/combo/predict.py b/combo/predict.py
index ff4d72b..50babf4 100644
--- a/combo/predict.py
+++ b/combo/predict.py
@@ -194,14 +194,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"),
diff --git a/scripts/predict_iwpt21.py b/scripts/predict_iwpt21.py
index 513ce3b..61b9cf7 100644
--- a/scripts/predict_iwpt21.py
+++ b/scripts/predict_iwpt21.py
@@ -36,6 +36,8 @@ flags.DEFINE_integer(name="cuda_device", default=-1,
help="Cuda device id (-1 for cpu).")
flags.DEFINE_boolean(name="expect_prefix", default=True,
help="Whether to expect allennlp prefix.")
+flags.DEFINE_integer(name="batch_size", default=32,
+ help="Batch size.")
def run(_):
@@ -68,6 +70,7 @@ def run(_):
--input_file {test_file}
--output_file {output_pred}
--cuda_device {FLAGS.cuda_device}
+ --batch_size {FLAGS.batch_size}
--silent
"""
utils.execute_command(command)
--
GitLab