Repository: awslabs/sockeye

Issue:
Bug in lexically-constrained decoding
It seems that my sampling commit (094baca) broke lexically constrained decoding. Here are BLEU scores against a random 4-word phrase from the reference:

- cb441b6 45.7 current master
- 094baca 45.7 sampling added. this probably broke everything (confirmed)
- b19685e 48.6 just before adding sampling. scores are better.

Also, decoding with constraints with a batch size of > 2 is broken.

I will fix both of these.


Discussion:

Resolving patch:
diff --git a/sockeye/__init__.py b/sockeye/__init__.py
--- a/sockeye/__init__.py
+++ b/sockeye/__init__.py
@@ -11,4 +11,4 @@
 # express or implied. See the License for the specific language governing
 # permissions and limitations under the License.
 
-__version__ = '1.18.83'
+__version__ = '1.18.84'
diff --git a/sockeye/lexical_constraints.py b/sockeye/lexical_constraints.py
--- a/sockeye/lexical_constraints.py
+++ b/sockeye/lexical_constraints.py
@@ -510,21 +510,18 @@ def topk(timestep: int,
         the updated constrained hypotheses, and the updated set of inactive hypotheses.
     """
 
-    if timestep == 1:
-        beam_size = 1
-
     for sentno in range(batch_size):
-        rows = slice(sentno * beam_size, (sentno + 1) * beam_size)
+        rows = slice(sentno * beam_size, sentno * beam_size + beam_size)
         if hypotheses[rows.start] is not None and hypotheses[rows.start].size() > 0:
             best_ids[rows], best_word_ids[rows], seq_scores[rows], \
-            hypotheses[rows], inactive[rows] = _topk(timestep,
-                                                     beam_size,
-                                                     inactive[rows],
-                                                     scores[rows],
-                                                     hypotheses[rows],
-                                                     best_ids[rows] - rows.start,
-                                                     best_word_ids[rows],
-                                                     seq_scores[rows])
+                hypotheses[rows], inactive[rows] = _sequential_topk(timestep,
+                                                                    beam_size,
+                                                                    inactive[rows],
+                                                                    scores[rows],
+                                                                    hypotheses[rows],
+                                                                    best_ids[rows] - rows.start,
+                                                                    best_word_ids[rows],
+                                                                    seq_scores[rows])
 
             # offsetting since the returned smallest_k() indices were slice-relative
             best_ids[rows] += rows.start
@@ -536,15 +533,15 @@ def topk(timestep: int,
     return best_ids, best_word_ids, seq_scores, hypotheses, inactive
 
 
-def _topk(timestep: int,
-          beam_size: int,
-          inactive: mx.nd.NDArray,
-          scores: mx.nd.NDArray,
-          hypotheses: List[ConstrainedHypothesis],
-          best_ids: mx.nd.NDArray,
-          best_word_ids: mx.nd.NDArray,
-          sequence_scores: mx.nd.NDArray) -> Tuple[np.array, np.array, np.array,
-                                                   List[ConstrainedHypothesis], mx.nd.NDArray]:
+def _sequential_topk(timestep: int,
+                     beam_size: int,
+                     inactive: mx.nd.NDArray,
+                     scores: mx.nd.NDArray,
+                     hypotheses: List[ConstrainedHypothesis],
+                     best_ids: mx.nd.NDArray,
+                     best_word_ids: mx.nd.NDArray,
+                     sequence_scores: mx.nd.NDArray) -> Tuple[np.array, np.array, np.array,
+                                                              List[ConstrainedHypothesis], mx.nd.NDArray]:
     """
     Builds a new topk list such that the beam contains hypotheses having completed different numbers of constraints.
     These items are built from three different types: (1) the best items across the whole
@@ -569,7 +566,7 @@ def _topk(timestep: int,
     for row, col, seq_score in zip(best_ids, best_word_ids, sequence_scores):
         row = int(row.asscalar())
         col = int(col.asscalar())
-        if hypotheses[row].is_valid(col):
+        if hypotheses[row] is not None and hypotheses[row].is_valid(col):
             seq_score = float(seq_score.asscalar())
             new_item = hypotheses[row].advance(col)
             cand = ConstrainedCandidate(row, col, seq_score, new_item)
@@ -579,7 +576,7 @@ def _topk(timestep: int,
     # (3) the best item (constrained or not) in that row
     best_next = mx.nd.argmin(scores, axis=1)
     for row in range(beam_size):
-        if inactive[row] or (timestep == 1 and row > 0):
+        if inactive[row]:
             continue
 
         hyp = hypotheses[row]


Test patch:
diff --git a/test/common.py b/test/common.py
--- a/test/common.py
+++ b/test/common.py
@@ -401,9 +401,7 @@ def test_translate_equivalence(data: Dict[str, Any], translate_params_equiv: str
     assert 'test_outputs' in data and 'test_scores' in data
     assert all(a == b for a, b in zip(data['test_outputs'], translate_outputs_equiv))
     assert all(abs(a - b) < 0.01 or np.isnan(a - b) for a, b in zip(data['test_scores'], translate_scores_equiv))
-
-
-def _create_constrained_inputs(translate_inputs: List[str], translate_outputs: List[str]) -> List[Dict[str, Any]]:
+def _create_reference_constraints(translate_inputs: List[str], translate_outputs: List[str]) -> List[Dict[str, Any]]:
     constrained_inputs = []
     for sentno, (source, translate_output) in enumerate(zip(translate_inputs, translate_outputs)):
         constrained_inputs.append(json.dumps({'text': source, 'constraints': ['
 {} 
'.format(translate_output)]}, ensure_ascii=False))
@@ -411,7 +409,7 @@ def _create_constrained_inputs(translate_inputs: List[str], translate_outputs: L
 
 
 def test_constrained_decoding_against_ref(data: Dict[str, Any], translate_params: str):
-    constrained_inputs = _create_constrained_inputs(data['test_inputs'], data['test_outputs'])
+    constrained_inputs = _create_reference_constraints(data['test_inputs'], data['test_outputs'])
     new_test_source_path = os.path.join(data['work_dir'], "test_constrained.txt")
     with open(new_test_source_path, 'w') as out:
         for json_line in constrained_inputs:
diff --git a/test/integration/test_constraints_int.py b/test/integration/test_constraints_int.py
--- a/test/integration/test_constraints_int.py
+++ b/test/integration/test_constraints_int.py
@@ -42,7 +42,7 @@
      " --rnn-attention-num-hidden 8 --loss cross-entropy --optimized-metric perplexity --max-updates 2"
      " --checkpoint-interval 2 --optimizer adam --initial-learning-rate 0.01 --batch-type sentence "
      " --decode-and-evaluate 0",
-     "--batch-size 2 --beam-size 10"),
+     "--batch-size 3 --beam-size 10 --beam-prune 1"),
     # Full transformer
     ("--encoder transformer --decoder transformer"
      " --num-layers 2 --transformer-attention-heads 2 --transformer-model-size 8 --num-embed 8"