Work around a TensorFlow bug when only one answer is extracted

C.f. tensorflow/nmt#117
Maluuba · May 25, 2018 · 0001f72 · 0001f72
1 parent 707087e
commit 0001f72
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/qgen/model.py b/qgen/model.py
@@ -132,7 +132,14 @@
 
 batch = expand_answers(batch, answers)
 
-helper = seq2seq.GreedyEmbeddingHelper(embedding, tf.fill([batch["size"]], START_TOKEN), END_TOKEN)
+# Work around https://github.com/tensorflow/nmt/issues/117
+class FixedHelper(seq2seq.GreedyEmbeddingHelper):
+    def sample(self, *args, **kwargs):
+        result = super().sample(*args, **kwargs)
+        result.set_shape([batch["size"]])
+        return result
+
+helper = FixedHelper(embedding, tf.fill([batch["size"]], START_TOKEN), END_TOKEN)
 decoder = seq2seq.BasicDecoder(decoder_cell, helper, encoder_state, output_layer=projection)
 
 if batch["size"] > 0: