Raise an error if n_best is greater than the beam width
@@ -173,6 +173,9 @@ class SequenceToSequence(Model):
n_best = params and params.get("n_best")
n_best = n_best or 1
+ if n_best > len(prediction["tokens"]):
+ raise ValueError("n_best cannot be greater than beam_width")
+
all_preds = []
for i in range(n_best):
tokens = prediction["tokens"][i][:prediction["length"][i] - 1] # Ignore </s>.
