feat: add beam_search_v2

gongel · gongel · commit 18b9ed600524 · 2021-08-03T14:45:13.000+08:00
diff --git a/examples/machine_translation/transformer/predict_beamsearch_v2.py b/examples/machine_translation/transformer/predict_beamsearch_v2.py
@@ -0,0 +1,116 @@
+import os
+import yaml
+import argparse
+from pprint import pprint
+from attrdict import AttrDict
+
+import paddle
+from paddlenlp.transformers import TransformerModel, position_encoding_init
+import reader
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--config",
+        default="./configs/transformer.base.yaml",
+        type=str,
+        help="Path of the config file. ")
+    args = parser.parse_args()
+    return args
+
+
+def post_process_seq(seq, bos_idx, eos_idx, output_bos=False, output_eos=False):
+    """
+    Post-process the decoded sequence.
+    """
+    eos_pos = len(seq) - 1
+    for i, idx in enumerate(seq):
+        if idx == eos_idx:
+            eos_pos = i
+            break
+    seq = [
+        idx for idx in seq[:eos_pos + 1]
+        if (output_bos or idx != bos_idx) and (output_eos or idx != eos_idx)
+    ]
+    return seq
+
+
+def do_predict(args):
+    if args.device == "gpu":
+        place = "gpu"
+    else:
+        place = "cpu"
+
+    paddle.set_device(place)
+
+    # Define data loader
+    test_loader, to_tokens = reader.create_infer_loader(args)
+
+    # Define model
+    # `TransformerGenerator` automatically chioces using `FasterTransformer`
+    # (with jit building) or the slower verison `InferTransformerModel`.
+    transformer = TransformerModel(
+        src_vocab_size=args.src_vocab_size,
+        trg_vocab_size=args.trg_vocab_size,
+        max_length=args.max_length + 1,
+        num_encoder_layers=args.n_layer,
+        num_decoder_layers=args.n_layer,
+        n_head=args.n_head,
+        d_model=args.d_model,
+        d_inner_hid=args.d_inner_hid,
+        dropout=args.dropout,
+        weight_sharing=args.weight_sharing,
+        bos_id=args.bos_idx,
+        eos_id=args.eos_idx)
+
+    # Load the trained model
+    assert args.init_from_params, (
+        "Please set init_from_params to load the infer model.")
+    model_dict = paddle.load(
+        os.path.join(args.init_from_params, "transformer.pdparams"))
+
+    # To avoid a longer length than training, reset the size of position
+    # encoding to max_length
+    model_dict["src_pos_embedding.pos_encoder.weight"] = position_encoding_init(
+        args.max_length + 1, args.d_model)
+    model_dict["trg_pos_embedding.pos_encoder.weight"] = position_encoding_init(
+        args.max_length + 1, args.d_model)
+
+    # Load the model_dict
+    transformer.load_dict(model_dict)
+
+    # Set evaluate mode
+    transformer.eval()
+
+    f = open(args.output_file, "w", encoding="utf-8")
+
+    with paddle.no_grad():
+        for (src_word, ) in test_loader:
+            # The shape of finished_seq is `[seq_len, batch_size, beam_size]`
+            # when `output_time_major` argument is `True` for TransformerGenerator.
+            finished_seq, finished_scores = transformer.beam_search_v2(
+                src_word=src_word,
+                beam_size=args.beam_size,
+                max_len=args.max_out_len,
+                alpha=0.6)
+            finished_seq = finished_seq.numpy()
+            for ins in finished_seq:
+                for beam_idx, beam in enumerate(ins):
+                    if beam_idx >= args.n_best:
+                        break
+                    id_list = post_process_seq(beam, args.bos_idx, args.eos_idx)
+                    word_list = to_tokens(id_list)
+                    sequence = " ".join(word_list) + "\n"
+                    f.write(sequence)
+    f.close()
+
+
+if __name__ == "__main__":
+    ARGS = parse_args()
+    yaml_file = ARGS.config
+    with open(yaml_file, 'rt') as f:
+        args = AttrDict(yaml.safe_load(f))
+    pprint(args)
+
+    do_predict(args)
diff --git a/paddlenlp/transformers/transformer/modeling.py b/paddlenlp/transformers/transformer/modeling.py
@@ -785,6 +785,249 @@ def forward(self, src_word, trg_word):
 
         return predict
 
+    def beam_search_v2(self, src_word, beam_size=4, max_len=None, alpha=0.6):
+        """
+        Beam search with the alive and finished two queues, both have a beam size
+        capicity separately. It includes `grow_topk` `grow_alive` `grow_finish` as
+        steps.
+        1. `grow_topk` selects the top `2*beam_size` candidates to avoid all getting
+        EOS.
+        2. `grow_alive` selects the top `beam_size` non-EOS candidates as the inputs
+        of next decoding step.
+        3. `grow_finish` compares the already finished candidates in the finished queue
+        and newly added finished candidates from `grow_topk`, and selects the top
+        `beam_size` finished candidates.
+        """
+
+        def expand_to_beam_size(tensor, beam_size):
+            tensor = paddle.reshape(tensor,
+                                    [tensor.shape[0], 1] + tensor.shape[1:])
+            tile_dims = [1] * len(tensor.shape)
+            tile_dims[1] = beam_size
+            return paddle.tile(tensor, tile_dims)
+
+        def merge_beam_dim(tensor):
+            return paddle.reshape(tensor, [-1] + tensor.shape[2:])
+
+        # run encoder
+        src_max_len = paddle.shape(src_word)[-1]
+        src_slf_attn_bias = paddle.cast(
+            src_word == self.bos_id,
+            dtype=paddle.get_default_dtype()).unsqueeze([1, 2]) * -1e9
+        src_slf_attn_bias.stop_gradient = True
+        src_pos = paddle.cast(
+            src_word != self.bos_id, dtype="int64") * paddle.arange(
+                start=0, end=src_max_len)
+        src_emb = self.src_word_embedding(src_word)
+        src_pos_emb = self.src_pos_embedding(src_pos)
+        src_emb = src_emb + src_pos_emb
+        enc_input = F.dropout(
+            src_emb, p=self.dropout,
+            training=self.training) if self.dropout else src_emb
+
+        enc_output = self.transformer.encoder(enc_input, src_slf_attn_bias)
+
+        # constant number
+        inf = float(1. * 1e7)
+        batch_size = enc_output.shape[0]
+        max_len = (enc_output.shape[1] + 20) if max_len is None else max_len
+
+        ### initialize states of beam search ###
+        ## init for the alive ##
+        initial_log_probs = paddle.to_tensor(
+            np.array(
+                [[0.] + [-inf] * (beam_size - 1)], dtype="float32"))
+        alive_log_probs = paddle.tile(initial_log_probs, [batch_size, 1])
+        alive_seq = paddle.to_tensor(
+            np.tile(
+                np.array(
+                    [[[self.bos_id]]], dtype="int64"), (batch_size, beam_size, 1
+                                                        )))
+
+        ## init for the finished ##
+        finished_scores = paddle.to_tensor(
+            np.array(
+                [[-inf] * beam_size], dtype="float32"))
+        finished_scores = paddle.tile(finished_scores, [batch_size, 1])
+        finished_seq = paddle.to_tensor(
+            np.tile(
+                np.array(
+                    [[[self.bos_id]]], dtype="int64"), (batch_size, beam_size, 1
+                                                        )))
+        finished_flags = paddle.zeros_like(finished_scores)
+
+        ### initialize inputs and states of transformer decoder ###
+        ## init inputs for decoder, shaped `[batch_size*beam_size, ...]`
+        trg_word = paddle.reshape(alive_seq[:, :, -1],
+                                  [batch_size * beam_size, 1])
+        trg_src_attn_bias = src_slf_attn_bias
+        trg_src_attn_bias = merge_beam_dim(
+            expand_to_beam_size(trg_src_attn_bias, beam_size))
+        enc_output = merge_beam_dim(expand_to_beam_size(enc_output, beam_size))
+
+        ## init states (caches) for transformer, need to be updated according to selected beam
+        caches = self.transformer.decoder.gen_cache(enc_output, do_zip=False)
+
+        def update_states(caches, beam_idx, beam_size):
+            new_caches = []
+            for cache in caches:
+                k = gather_2d_by_gather(cache[0].k, beam_idx, beam_size,
+                                        batch_size, False)
+                v = gather_2d_by_gather(cache[0].v, beam_idx, beam_size,
+                                        batch_size, False)
+                new_caches.append((nn.MultiHeadAttention.Cache(k, v), cache[1]))
+            return new_caches
+
+        def gather_2d_by_gather(tensor_nd,
+                                beam_idx,
+                                beam_size,
+                                batch_size,
+                                need_flat=True):
+            batch_idx = paddle.arange(
+                0, batch_size, 1, dtype="int64") * beam_size
+            flat_tensor = merge_beam_dim(tensor_nd) if need_flat else tensor_nd
+            idx = paddle.reshape(
+                paddle.add(beam_idx, batch_idx.unsqueeze(-1)), [-1])
+            new_flat_tensor = paddle.gather(flat_tensor, idx)
+            new_tensor_nd = paddle.reshape(
+                new_flat_tensor,
+                shape=[batch_size, beam_idx.shape[1]] +
+                tensor_nd.shape[2:]) if need_flat else new_flat_tensor
+            return new_tensor_nd
+
+        def early_finish(alive_log_probs, finished_scores,
+                         finished_in_finished):
+            max_length_penalty = np.power(((5. + max_len) / 6.), alpha)
+            # The best possible score of the most likely alive sequence
+            lower_bound_alive_scores = alive_log_probs[:,
+                                                       0] / max_length_penalty
+
+            # Now to compute the lowest score of a finished sequence in finished
+            # If the sequence isn't finished, we multiply it's score by 0. since
+            # scores are all -ve, taking the min will give us the score of the lowest
+            # finished item.
+            lowest_score_of_fininshed_in_finished = paddle.min(
+                finished_scores * finished_in_finished, 1)
+            # If none of the sequences have finished, then the min will be 0 and
+            # we have to replace it by -ve INF if it is. The score of any seq in alive
+            # will be much higher than -ve INF and the termination condition will not
+            # be met.
+            lowest_score_of_fininshed_in_finished += (
+                1. - paddle.max(finished_in_finished, 1)) * -inf
+            bound_is_met = paddle.all(
+                paddle.greater_than(lowest_score_of_fininshed_in_finished,
+                                    lower_bound_alive_scores))
+
+            return bound_is_met
+
+        def grow_topk(i, logits, alive_seq, alive_log_probs, states):
+            logits = paddle.reshape(logits, [batch_size, beam_size, -1])
+            candidate_log_probs = paddle.log(F.softmax(logits, axis=2))
+            log_probs = paddle.add(candidate_log_probs,
+                                   alive_log_probs.unsqueeze(-1))
+
+            length_penalty = np.power(5.0 + (i + 1.0) / 6.0, alpha)
+            curr_scores = log_probs / length_penalty
+            flat_curr_scores = paddle.reshape(curr_scores, [batch_size, -1])
+
+            topk_scores, topk_ids = paddle.topk(
+                flat_curr_scores, k=beam_size * 2)
+
+            topk_log_probs = topk_scores * length_penalty
+
+            topk_beam_index = topk_ids // self.trg_vocab_size
+            topk_ids = topk_ids % self.trg_vocab_size
+
+            # use gather as gather_nd, TODO: use gather_nd
+            topk_seq = gather_2d_by_gather(alive_seq, topk_beam_index,
+                                           beam_size, batch_size)
+            topk_seq = paddle.concat(
+                [topk_seq, paddle.reshape(topk_ids, topk_ids.shape + [1])],
+                axis=2)
+            states = update_states(states, topk_beam_index, beam_size)
+            eos = paddle.full(
+                shape=topk_ids.shape, dtype="int64", fill_value=self.eos_id)
+            topk_finished = paddle.cast(paddle.equal(topk_ids, eos), "float32")
+
+            # topk_seq: [batch_size, 2*beam_size, i+1]
+            # topk_log_probs, topk_scores, topk_finished: [batch_size, 2*beam_size]
+            return topk_seq, topk_log_probs, topk_scores, topk_finished, states
+
+        def grow_alive(curr_seq, curr_scores, curr_log_probs, curr_finished,
+                       states):
+            curr_scores += curr_finished * -inf
+            _, topk_indexes = paddle.topk(curr_scores, k=beam_size)
+            alive_seq = gather_2d_by_gather(curr_seq, topk_indexes,
+                                            beam_size * 2, batch_size)
+            alive_log_probs = gather_2d_by_gather(curr_log_probs, topk_indexes,
+                                                  beam_size * 2, batch_size)
+            states = update_states(states, topk_indexes, beam_size * 2)
+
+            return alive_seq, alive_log_probs, states
+
+        def grow_finished(finished_seq, finished_scores, finished_flags,
+                          curr_seq, curr_scores, curr_finished):
+            # finished scores
+            finished_seq = paddle.concat(
+                [
+                    finished_seq, paddle.full(
+                        shape=[batch_size, beam_size, 1],
+                        dtype="int64",
+                        fill_value=self.eos_id)
+                ],
+                axis=2)
+            # Set the scores of the unfinished seq in curr_seq to large negative
+            # values
+            curr_scores += (1. - curr_finished) * -inf
+            # concatenating the sequences and scores along beam axis
+            curr_finished_seq = paddle.concat([finished_seq, curr_seq], axis=1)
+            curr_finished_scores = paddle.concat(
+                [finished_scores, curr_scores], axis=1)
+            curr_finished_flags = paddle.concat(
+                [finished_flags, curr_finished], axis=1)
+            _, topk_indexes = paddle.topk(curr_finished_scores, k=beam_size)
+            finished_seq = gather_2d_by_gather(curr_finished_seq, topk_indexes,
+                                               beam_size * 3, batch_size)
+            finished_scores = gather_2d_by_gather(
+                curr_finished_scores, topk_indexes, beam_size * 3, batch_size)
+            finished_flags = gather_2d_by_gather(
+                curr_finished_flags, topk_indexes, beam_size * 3, batch_size)
+            return finished_seq, finished_scores, finished_flags
+
+        for i in range(max_len):
+            trg_pos = paddle.full(
+                shape=trg_word.shape, dtype="int64", fill_value=i)
+            trg_emb = self.trg_word_embedding(trg_word)
+            trg_pos_emb = self.trg_pos_embedding(trg_pos)
+            trg_emb = trg_emb + trg_pos_emb
+            dec_input = F.dropout(
+                trg_emb, p=self.dropout,
+                training=self.training) if self.dropout else trg_emb
+
+            logits, caches = self.transformer.decoder(
+                dec_input, enc_output, None, trg_src_attn_bias, caches)
+            logits = paddle.reshape(
+                logits,
+                shape=[-1, logits.shape[-1]], )
+            logits = self.linear(logits)
+
+            topk_seq, topk_log_probs, topk_scores, topk_finished, states = grow_topk(
+                i, logits, alive_seq, alive_log_probs, caches)
+            alive_seq, alive_log_probs, states = grow_alive(
+                topk_seq, topk_scores, topk_log_probs, topk_finished, states)
+            caches = states
+            finished_seq, finished_scores, finished_flags = grow_finished(
+                finished_seq, finished_scores, finished_flags, topk_seq,
+                topk_scores, topk_finished)
+            trg_word = paddle.reshape(alive_seq[:, :, -1],
+                                      [batch_size * beam_size, 1])
+
+            if early_finish(alive_log_probs, finished_scores,
+                            finished_flags).numpy():
+                break
+
+        return finished_seq, finished_scores
+
 
 class InferTransformerModel(TransformerModel):
     """
