[bug] Beam Search Returns the Least Likely Sentence at the End

Enter the chapter number

Chapter 14, Beam Search

Enter the page number

No response

What is the cell's number in the notebook

No response

Enter the environment you are using to run the notebook

None

Describe your issue

Current Beam search implementation keeps track of the k most likely sentences, but then returns the least likely of the final k at the end:

def beam_search(model, src_text, beam_width=3, max_length=20,
                verbose=False, length_penalty=0.6):
    top_translations = [(torch.tensor(0.), "")]
    for index in range(max_length):
        if verbose:
            print(f"Top {beam_width} translations so far:")
            for log_proba, tgt_text in top_translations:
                print(f"    {log_proba.item():.3f} – {tgt_text}")

        candidates = []
        for log_proba, tgt_text in top_translations:
            if tgt_text.endswith(" </s>"):
                candidates.append((log_proba, tgt_text))
                continue  # don't add tokens after EOS token
            batch, _ = nmt_collate_fn([{"source_text": src_text,
                                        "target_text": tgt_text}])
            with torch.no_grad():
                Y_logits = model(batch.to(device))
                Y_log_proba = F.log_softmax(Y_logits, dim=1)
                Y_top_log_probas = torch.topk(Y_log_proba, k=beam_width, dim=1)

            for beam_index in range(beam_width):
                next_token_log_proba = Y_top_log_probas.values[0, beam_index, index]
                next_token_id = Y_top_log_probas.indices[0, beam_index, index]
                next_token = nmt_tokenizer.id_to_token(next_token_id)
                next_tgt_text = tgt_text + " " + next_token
                candidates.append((log_proba + next_token_log_proba, next_tgt_text))

        def length_penalized_score(candidate, alpha=length_penalty):
            log_proba, text = candidate
            length = len(text.split())
            penalty = ((5 + length) ** alpha) / (6 ** alpha)
            return log_proba / penalty

        top_translations = sorted(candidates,
                                  key=length_penalized_score,
                                  reverse=True)[:beam_width]

    return top_translations[-1][1]

Enter what you expected to happen

No response

If you found a workaround, describe it here

Last line should be:

return top_translations[0][1]