def test_t5_forward_beam_search(model: T5, tokenizer: hf_t5.T5Tokenizer):
    """
    Make sure beam search generates reasonable results, and that we get the same results
    for a given input, regardless of whether we run it on its own or part of a batch.
    """

    def run_beam_search(sents):
        input_ids = tokenizer(sents, return_tensors="pt", padding=True).input_ids
        outputs = model(input_ids)
        preds = [tokenizer.decode(preds[0]) for preds in outputs.predictions]
        probs = [round(float(p[0]), 4) for p in outputs.predicted_log_probs.exp()]
        return preds, probs

    s1 = "translate English to German: That is good"
    s2 = (
        "mnli premise: The Old One always comforted Ca'daan, except today. "
        "hypothesis: Ca'daan knew the Old One very well."
    )

    s1_pred, s1_prob = run_beam_search([s1])
    assert s1_pred == ["Das ist gut.</s>"]
    assert s1_prob == [0.5645]

    s2_pred, s2_prob = run_beam_search([s2])
    assert s2_pred == ["neutral</s> </s> </s> </s>"]
    assert s2_prob == [0.3992]

    combined_preds, combined_probs = run_beam_search([s1, s2])
    assert combined_preds == s1_pred + s2_pred
    assert combined_probs == s1_prob + s2_prob
Make sure beam search generates reasonable results, and that we get the same results
for a given input, regardless of whether we run it on its own or part of a batch.
tests/modules/transformer/t5_test.py