adding search.PrefixConstrainedBeamSearch

fairseq/search.py

@@ -124,6 +124,58 @@ def step(self, step: int, lprobs, scores: Optional[Tensor]):
         # At this point, beams_buf and indices_buf are single-dim and contain relative indices
         return scores_buf, indices_buf, beams_buf
 
+
+class PrefixConstrainedBeamSearch(Search):
+    def __init__(self, tgt_dict, prefix_allowed_tokens_fn):
+        super().__init__(tgt_dict)
+        self.prefix_allowed_tokens_fn = prefix_allowed_tokens_fn
+
+    @torch.jit.export
+    def apply_mask(self, x, prev_output_tokens, original_batch_idxs):
+        beam_size = x.shape[0] // original_batch_idxs.shape[0]
+        original_batch_idxs = (
+            original_batch_idxs.unsqueeze(-1).repeat((1, beam_size)).flatten().tolist()
+        )
+
+        mask = torch.full_like(x, -math.inf)
+        for sent_i, (sent, batch_i) in enumerate(
+            zip(prev_output_tokens, original_batch_idxs)
+        ):
+            mask[sent_i, :, self.prefix_allowed_tokens_fn(batch_i, sent)] = 0
+
+        return mask
+
+    @torch.jit.export
+    def step(self, step: int, lprobs, scores: Tensor, prev_output_tokens: Tensor, original_batch_idxs: Tensor):
+        bsz, beam_size, vocab_size = lprobs.size()
+
+        lprobs += self.apply_mask(lprobs.view(bsz * beam_size, 1, vocab_size), prev_output_tokens,
+                                  original_batch_idxs).view(bsz, beam_size, vocab_size)
+
+        if step == 0:
+            # at the first step all hypotheses are equally likely, so use
+            # only the first beam
+            lprobs = lprobs[:, ::beam_size, :].contiguous()
+        else:
+            # make probs contain cumulative scores for each hypothesis
+            assert scores is not None
+            lprobs = lprobs + scores[:, :, step - 1].unsqueeze(-1)
+
+        top_prediction = torch.topk(
+            lprobs.view(bsz, -1),
+            k=min(
+                # Take the best beam_size predictions. We'll choose the first
+                # beam_size of these which don't predict eos to continue with.
+                beam_size,
+                lprobs.view(bsz, -1).size(1) - 1,  # -1 so we never select pad
+            ),
+        )
+        scores_buf = top_prediction[0]
+        indices_buf = top_prediction[1]
+        beams_buf = indices_buf // vocab_size
+        indices_buf = indices_buf.fmod(vocab_size)
+        return scores_buf, indices_buf, beams_buf
+
 
 class LexicallyConstrainedBeamSearch(Search):
     """Implements lexically constrained beam search as described in

fairseq/sequence_generator.py

@@ -269,6 +269,7 @@ def _generate(
 
         reorder_state: Optional[Tensor] = None
         batch_idxs: Optional[Tensor] = None
+        original_batch_idxs = sample.get("id", torch.arange(0, bsz))
         for step in range(max_len + 1):  # one extra step for EOS marker
             # reorder decoder internal states based on the prev choice of beams
             # print(f'step: {step}')
@@ -281,6 +282,7 @@ def _generate(
                     reorder_state.view(-1, beam_size).add_(
                         corr.unsqueeze(-1) * beam_size
                     )
+                    original_batch_idxs = original_batch_idxs[batch_idxs]
                 self.model.reorder_incremental_state(incremental_states, reorder_state)
                 encoder_outs = self.model.reorder_encoder_out(
                     encoder_outs, reorder_state
@@ -338,12 +340,21 @@ def _generate(
                 lprobs = self._no_repeat_ngram(tokens, lprobs, bsz, beam_size, step)
 
             # Shape: (batch, cand_size)
-            cand_scores, cand_indices, cand_beams = self.search.step(
-                step,
-                lprobs.view(bsz, -1, self.vocab_size),
-                scores.view(bsz, beam_size, -1)[:, :, :step],
-            )
-
+            if isinstance(self.search, search.PrefixConstrainedBeamSearch):
+                cand_scores, cand_indices, cand_beams = self.search.step(
+                    step,
+                    lprobs.view(bsz, -1, self.vocab_size),
+                    scores.view(bsz, beam_size, -1)[:, :, :step],
+                    tokens[:, : step + 1],
+                    original_batch_idxs,
+                )
+            else:
+                cand_scores, cand_indices, cand_beams = self.search.step(
+                    step,
+                    lprobs.view(bsz, -1, self.vocab_size),
+                    scores.view(bsz, beam_size, -1)[:, :, :step],
+                )
+
             # cand_bbsz_idx contains beam indices for the top candidate
             # hypotheses, with a range of values: [0, bsz*beam_size),
             # and dimensions: [bsz, cand_size]
@@ -385,8 +396,10 @@ def _generate(
             assert num_remaining_sent >= 0
             if num_remaining_sent == 0:
                 break
+            if isinstance(self.search, search.PrefixConstrainedBeamSearch) and step >= max_len:
+                break
             assert step < max_len
-
+            
             # Remove finalized sentences (ones for which {beam_size}
             # finished hypotheses have been generated) from the batch.
             if len(finalized_sents) > 0:

fairseq/tasks/fairseq_task.py

@@ -312,6 +312,7 @@ def build_generator(
         match_source_len = getattr(args, "match_source_len", False)
         diversity_rate = getattr(args, "diversity_rate", -1)
         constrained = getattr(args, "constraints", False)
+        prefix_allowed_tokens_fn = getattr(args, "prefix_allowed_tokens_fn", None)
         if (
             sum(
                 int(cond)
@@ -353,6 +354,8 @@ def build_generator(
             )
         elif constrained:
             search_strategy = search.LexicallyConstrainedBeamSearch(self.target_dictionary, args.constraints)
+        elif prefix_allowed_tokens_fn:
+            search_strategy = search.PrefixConstrainedBeamSearch(self.target_dictionary, prefix_allowed_tokens_fn)
         else:
             search_strategy = search.BeamSearch(self.target_dictionary)