⚡️ Speed up method BeamSearchDecoder.finalize by 41%

[codeflash-ai]

📄 41% (0.41x) speedup for BeamSearchDecoder.finalize in whisper/decoding.py

⏱️ Runtime : 2.66 milliseconds → 1.88 milliseconds (best of 570 runs)

📝 Explanation and details

Here's a faster version of your code.
Main changes.

• Replace np.argsort with PyTorch's torch.argsort to avoid unnecessary conversion between numpy arrays and torch tensors (which can be expensive on GPU).
• Avoid repeated .tolist() and multiple conversions between data types.
• Pre-allocate lists with list comprehensions instead of updating dicts with tuple keys (where not required).
• Simplify loops to reduce the number of Python-side operations.
• Remove unnecessary variables and avoid repeated lookups.

Key optimization details:

• Uses torch's topk for quick sorting and indexing instead of numpy.
• Reduces nested for loops and conversions.
• Keeps your logic and signatures identical.

If you must have unique sequence keys (as before) and the sequences could possibly repeat, the original dict usage is retained. However, for even more speed, consider using a list if uniqueness is not critical or cannot be violated in your context.

Let me know if your finished_sequences can be a list-of-lists and you're only interested in speed at the expense of removing uniqueness enforcement! This version, though, already avoids numpy and should be considerably faster in batched GPU/CPU settings.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 25 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests Details

from typing import Dict, List, Optional, Tuple

import numpy as np
# imports
import pytest  # used for our unit tests
import torch
from torch import Tensor
from whisper.decoding import BeamSearchDecoder


# Dummy Inference and TokenDecoder classes for testability
class Inference:
    pass

class TokenDecoder:
    pass
from whisper.decoding import BeamSearchDecoder

# -----------------------------
# Unit tests for finalize
# -----------------------------

# Helper to create a dummy decoder with finished_sequences
def make_decoder(beam_size, eot, finished_sequences, patience=None):
    decoder = BeamSearchDecoder(beam_size, eot, Inference(), patience)
    decoder.finished_sequences = finished_sequences
    return decoder

# 1. BASIC TEST CASES

def test_finalize_basic_enough_finished_sequences():
    """
    Basic test: finished_sequences already has enough sequences.
    finalize should not add any new ones.
    """
    beam_size = 2
    eot = 99
    finished_sequences = [
        { (1,2,3,99): -1.0, (1,2,4,99): -2.0 },  # batch 0
        { (5,6,99): -0.5, (5,7,99): -1.5 }       # batch 1
    ]
    decoder = make_decoder(beam_size, eot, finished_sequences)
    # preceding_tokens and sum_logprobs are not used since all beams are finished
    preceding_tokens = torch.tensor([[[1,2,3],[1,2,4]], [[5,6],[5,7]]])
    sum_logprobs = torch.tensor([[0.0,0.0],[0.0,0.0]])
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)

def test_finalize_basic_not_enough_finished_sequences():
    """
    Basic test: finished_sequences has only one finished, beam_size=2, so one should be added.
    """
    beam_size = 2
    eot = 99
    finished_sequences = [
        { (1,2,3,99): -1.0 },  # batch 0, only 1 finished
    ]
    decoder = make_decoder(beam_size, eot, finished_sequences)
    # 2 beams, 2 logprobs
    preceding_tokens = torch.tensor([ [ [1,2,3], [1,2,4] ] ])  # shape (1,2,3)
    sum_logprobs = torch.tensor([ [ -1.0, -0.5 ] ])            # shape (1,2)
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)

def test_finalize_basic_all_unfinished():
    """
    Basic test: finished_sequences empty, should fill up with top beams.
    """
    beam_size = 2
    eot = 99
    finished_sequences = [
        {},  # batch 0
    ]
    decoder = make_decoder(beam_size, eot, finished_sequences)
    preceding_tokens = torch.tensor([ [ [1,2,3], [1,2,4] ] ])  # shape (1,2,3)
    sum_logprobs = torch.tensor([ [ -1.0, -0.5 ] ])            # shape (1,2)
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)

def test_finalize_basic_multiple_batches():
    """
    Basic test: batch size > 1, some finished, some not.
    """
    beam_size = 2
    eot = 99
    finished_sequences = [
        { (1,2,3,99): -1.0 },  # batch 0
        {},                    # batch 1
    ]
    decoder = make_decoder(beam_size, eot, finished_sequences)
    preceding_tokens = torch.tensor([
        [ [1,2,3], [1,2,4] ],
        [ [5,6], [5,7] ]
    ])  # shape (2,2,3)
    sum_logprobs = torch.tensor([
        [ -1.0, -0.5 ],
        [ -0.5, -1.5 ]
    ])
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)

# 2. EDGE TEST CASES

def test_finalize_edge_empty_finished_and_zero_beam():
    """
    Edge: beam_size=0 should assert
    """
    with pytest.raises(AssertionError):
        make_decoder(0, 99, [{}])

def test_finalize_edge_patience_zero():
    """
    Edge: patience=0 should assert
    """
    with pytest.raises(AssertionError):
        make_decoder(2, 99, [{}], patience=0)

def test_finalize_edge_patience_fractional():
    """
    Edge: patience fractional, should round beam_size*patience
    """
    beam_size = 3
    patience = 1.5
    decoder = make_decoder(beam_size, 99, [{}], patience=patience)

def test_finalize_edge_duplicate_sequences():
    """
    Edge: If unfinished sequence matches a finished one, should not duplicate.
    """
    beam_size = 2
    eot = 99
    finished_sequences = [
        { (1,2,3,99): -1.0 },  # already finished
    ]
    decoder = make_decoder(beam_size, eot, finished_sequences)
    # The unfinished beam is the same as finished
    preceding_tokens = torch.tensor([ [ [1,2,3], [1,2,3] ] ])
    sum_logprobs = torch.tensor([ [ -1.0, -0.5 ] ])
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)

def test_finalize_edge_sum_logprobs_ties():
    """
    Edge: sum_logprobs are tied, should still fill up beams.
    """
    beam_size = 2
    eot = 99
    finished_sequences = [ {} ]
    decoder = make_decoder(beam_size, eot, finished_sequences)
    preceding_tokens = torch.tensor([ [ [1,2,3], [1,2,4] ] ])
    sum_logprobs = torch.tensor([ [ -1.0, -1.0 ] ])
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)

def test_finalize_edge_preceding_tokens_shorter_than_eot():
    """
    Edge: preceding_tokens shorter than eot, should append eot.
    """
    beam_size = 1
    eot = 42
    finished_sequences = [ {} ]
    decoder = make_decoder(beam_size, eot, finished_sequences)
    preceding_tokens = torch.tensor([ [ [1] ] ])
    sum_logprobs = torch.tensor([ [ -2.0 ] ])
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)

def test_finalize_edge_preceding_tokens_empty():
    """
    Edge: preceding_tokens empty, should only append eot.
    """
    beam_size = 1
    eot = 7
    finished_sequences = [ {} ]
    decoder = make_decoder(beam_size, eot, finished_sequences)
    preceding_tokens = torch.zeros((1,1,0), dtype=torch.long)
    sum_logprobs = torch.tensor([ [ -3.0 ] ])
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)

def test_finalize_edge_large_patience_and_beam():
    """
    Edge: Large patience and beam_size, but not exceeding limits.
    """
    beam_size = 10
    patience = 2.0
    finished_sequences = [ {} ]
    decoder = make_decoder(beam_size, 99, finished_sequences, patience=patience)

# 3. LARGE SCALE TEST CASES

def test_finalize_large_batch_and_beams():
    """
    Large scale: batch size 10, beam size 10, all unfinished.
    """
    batch_size = 10
    beam_size = 10
    seq_len = 5
    eot = 100
    finished_sequences = [ {} for _ in range(batch_size) ]
    decoder = make_decoder(beam_size, eot, finished_sequences)
    # Each beam is [batch, beam, seq]
    preceding_tokens = torch.arange(batch_size * beam_size * seq_len).reshape(batch_size, beam_size, seq_len)
    # Logprobs: make the first beam highest, rest decreasing
    sum_logprobs = torch.stack([
        torch.linspace(-0.1, -beam_size, beam_size) for _ in range(batch_size)
    ])
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)
    for b in range(batch_size):
        # Should be sorted by logprob descending
        for k in range(beam_size):
            expected = preceding_tokens[b, k].tolist() + [eot]

def test_finalize_large_filled_and_unfilled():
    """
    Large scale: some finished, some not, batch size 5, beam 5.
    """
    batch_size = 5
    beam_size = 5
    seq_len = 3
    eot = 88
    finished_sequences = []
    # For even batches, all finished; for odd, none finished
    for i in range(batch_size):
        if i % 2 == 0:
            finished_sequences.append({ tuple([i,i,i,eot]): -i })
        else:
            finished_sequences.append({})
    decoder = make_decoder(beam_size, eot, finished_sequences)
    preceding_tokens = torch.randint(0, 10, (batch_size, beam_size, seq_len))
    sum_logprobs = torch.rand(batch_size, beam_size)
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)
    for i in range(batch_size):
        pass

def test_finalize_large_max_capacity():
    """
    Large scale: test max_candidates with patience > 1, but only beam_size are returned.
    """
    batch_size = 2
    beam_size = 8
    patience = 1.5
    seq_len = 4
    eot = 77
    finished_sequences = [ {} for _ in range(batch_size) ]
    decoder = make_decoder(beam_size, eot, finished_sequences, patience=patience)
    preceding_tokens = torch.randint(0, 20, (batch_size, beam_size, seq_len))
    sum_logprobs = torch.rand(batch_size, beam_size)
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)
    for b in range(batch_size):
        pass

def test_finalize_large_all_finished():
    """
    Large scale: all finished_sequences filled, nothing should be added.
    """
    batch_size = 3
    beam_size = 4
    eot = 42
    finished_sequences = []
    for i in range(batch_size):
        finished_sequences.append({
            tuple([i,i,i,eot]): -i,
            tuple([i,i,i+1,eot]): -i-1,
            tuple([i,i,i+2,eot]): -i-2,
            tuple([i,i,i+3,eot]): -i-3,
        })
    decoder = make_decoder(beam_size, eot, finished_sequences)
    preceding_tokens = torch.zeros((batch_size, beam_size, 3), dtype=torch.long)
    sum_logprobs = torch.zeros((batch_size, beam_size))
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)
    for b in range(batch_size):
        # Check all sequences end with eot
        for t in tokens[b]:
            pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from typing import Dict, List, Optional, Tuple

# imports
import pytest  # used for our unit tests
import torch
from torch import Tensor
from whisper.decoding import BeamSearchDecoder


# Dummy classes to allow BeamSearchDecoder to instantiate
class Inference:
    pass

class TokenDecoder:
    pass
from whisper.decoding import BeamSearchDecoder

# -------------------- UNIT TESTS --------------------

# Helper function to create dummy finished_sequences
def make_finished_sequences(batch_size, finished_per_batch):
    """
    Returns a list of dicts, each dict maps tuple(token_ids) -> score.
    finished_per_batch: list of lists of (token_seq, score)
    """
    out = []
    for seqs in finished_per_batch:
        d = {}
        for tokens, score in seqs:
            d[tuple(tokens)] = score
        out.append(d)
    return out

# ---- BASIC TEST CASES ----

def test_finalize_all_finished_exact_beam_size():
    """
    All sequences finished, each batch has exactly beam_size finished sequences.
    finalize should not add any new sequences.
    """
    beam_size = 2
    eot = 99
    decoder = BeamSearchDecoder(beam_size=beam_size, eot=eot, inference=Inference())
    # batch_size = 1
    finished_sequences = [
        [([1, 2, 3], 0.5), ([4, 5, 6], 0.8)]
    ]
    decoder.finished_sequences = make_finished_sequences(1, finished_sequences)
    # preceding_tokens and sum_logprobs should not matter
    preceding_tokens = torch.tensor([[[1, 2, 3], [4, 5, 6]]])
    sum_logprobs = torch.tensor([[0.5, 0.8]])
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)

def test_finalize_not_enough_finished_adds_best_unfinished():
    """
    Not enough finished sequences: finalize should add best scoring unfinished ones.
    """
    beam_size = 3
    eot = 88
    decoder = BeamSearchDecoder(beam_size=beam_size, eot=eot, inference=Inference())
    # batch_size = 1, only 1 finished
    finished_sequences = [
        [([10, 20], 0.4)]
    ]
    decoder.finished_sequences = make_finished_sequences(1, finished_sequences)
    # 2 unfinished candidates
    preceding_tokens = torch.tensor([[[1,2], [3,4]]])
    sum_logprobs = torch.tensor([[0.9, 0.7]])
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)
    all_seqs = [tuple(t.tolist()) for t in tokens[0]]
    # logprobs should match
    idx_10_20 = all_seqs.index((10,20))
    idx_1_2_88 = all_seqs.index((1,2,88))
    idx_3_4_88 = all_seqs.index((3,4,88))

def test_finalize_multiple_batches():
    """
    Multiple batches, each with different number of finished/unfinished.
    """
    beam_size = 2
    eot = 7
    decoder = BeamSearchDecoder(beam_size=beam_size, eot=eot, inference=Inference())
    finished_sequences = [
        [([1,2], 1.0)], # batch 0: 1 finished
        [([3,4], 2.0), ([5,6], 3.0)], # batch 1: 2 finished
    ]
    decoder.finished_sequences = make_finished_sequences(2, finished_sequences)
    # batch 0: 2 unfinished
    # batch 1: 1 unfinished
    preceding_tokens = torch.tensor([
        [[7,8],[9,10]],
        [[11,12],[13,14]]
    ])
    sum_logprobs = torch.tensor([
        [0.5, 0.6],
        [0.7, 0.8]
    ])
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)
    # batch 0: should include ([1,2], 1.0) and one of the unfinished with eot
    batch0_seqs = [tuple(t.tolist()) for t in tokens[0]]
    # batch 1: should be unchanged
    batch1_seqs = [tuple(t.tolist()) for t in tokens[1]]

# ---- EDGE TEST CASES ----

def test_finalize_no_finished_sequences():
    """
    No finished sequences at all: finalize should fill up all with best unfinished.
    """
    beam_size = 2
    eot = 42
    decoder = BeamSearchDecoder(beam_size=beam_size, eot=eot, inference=Inference())
    finished_sequences = [
        [] # batch 0
    ]
    decoder.finished_sequences = make_finished_sequences(1, finished_sequences)
    # 3 unfinished candidates, only top 2 should be chosen
    preceding_tokens = torch.tensor([[[1,2],[3,4],[5,6]]])
    sum_logprobs = torch.tensor([[0.2, 0.7, 0.5]])
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)
    seqs = [tuple(t.tolist()) for t in tokens[0]]
    # Logprobs should match
    idx_3_4_42 = seqs.index((3,4,42))
    idx_5_6_42 = seqs.index((5,6,42))

def test_finalize_duplicate_unfinished_sequences():
    """
    Unfinished candidates may produce duplicate sequences (by tokens+eot) as already finished.
    finalize should not add duplicate sequences.
    """
    beam_size = 2
    eot = 99
    decoder = BeamSearchDecoder(beam_size=beam_size, eot=eot, inference=Inference())
    finished_sequences = [
        [([1,2,99], 0.5)]
    ]
    decoder.finished_sequences = make_finished_sequences(1, finished_sequences)
    # Unfinished candidate will become [1,2,99] after appending eot
    preceding_tokens = torch.tensor([[[1,2],[3,4]]])
    sum_logprobs = torch.tensor([[0.9, 0.7]])
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)
    seqs = [tuple(t.tolist()) for t in tokens[0]]

def test_finalize_empty_batch():
    """
    Edge case: batch size 0 (no data).
    """
    beam_size = 2
    eot = 1
    decoder = BeamSearchDecoder(beam_size=beam_size, eot=eot, inference=Inference())
    finished_sequences = []
    decoder.finished_sequences = make_finished_sequences(0, finished_sequences)
    preceding_tokens = torch.empty((0, 2, 3), dtype=torch.long)
    sum_logprobs = torch.empty((0, 2))
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)

def test_finalize_patience_greater_than_1():
    """
    Patience > 1.0 should increase max_candidates, but finalize should still return beam_size per batch.
    """
    beam_size = 2
    eot = 55
    patience = 2.5
    decoder = BeamSearchDecoder(beam_size=beam_size, eot=eot, inference=Inference(), patience=patience)
    finished_sequences = [
        [([1,2], 0.1)]
    ]
    decoder.finished_sequences = make_finished_sequences(1, finished_sequences)
    preceding_tokens = torch.tensor([[[3,4],[5,6],[7,8],[9,10],[11,12]]])
    sum_logprobs = torch.tensor([[0.9, 0.8, 0.7, 0.6, 0.5]])
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)


def test_finalize_zero_beam_size_raises():
    """
    Invalid beam_size (<=0) should raise assertion error.
    """
    with pytest.raises(AssertionError):
        BeamSearchDecoder(beam_size=0, eot=0, inference=Inference())

# ---- LARGE SCALE TEST CASES ----

def test_finalize_large_batch_and_beam():
    """
    Large batch and beam size, all finished.
    """
    batch_size = 50
    beam_size = 10
    eot = 100
    decoder = BeamSearchDecoder(beam_size=beam_size, eot=eot, inference=Inference())
    # Each batch has beam_size finished sequences
    finished_sequences = []
    for i in range(batch_size):
        seqs = []
        for j in range(beam_size):
            seqs.append(([i, j], float(i + j)))
        finished_sequences.append(seqs)
    decoder.finished_sequences = make_finished_sequences(batch_size, finished_sequences)
    # preceding_tokens and sum_logprobs are irrelevant
    preceding_tokens = torch.zeros((batch_size, beam_size, 2), dtype=torch.long)
    sum_logprobs = torch.zeros((batch_size, beam_size))
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)
    for t in tokens:
        pass
    for l in logprobs:
        pass

def test_finalize_large_num_unfinished():
    """
    Large number of unfinished candidates, only top beam_size should be added.
    """
    batch_size = 3
    beam_size = 5
    eot = 77
    decoder = BeamSearchDecoder(beam_size=beam_size, eot=eot, inference=Inference())
    finished_sequences = [
        [], [], []
    ]
    decoder.finished_sequences = make_finished_sequences(batch_size, finished_sequences)
    # 20 unfinished per batch
    preceding_tokens = torch.arange(batch_size*20*3).reshape(batch_size, 20, 3)
    # Scores: sorted so last 5 are best
    sum_logprobs = torch.arange(batch_size*20).reshape(batch_size, 20).float()
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)
    for batch in range(batch_size):
        # The best 5 should be from indices 15-19
        for i in range(15, 20):
            seq = tuple(preceding_tokens[batch, i].tolist() + [eot])

def test_finalize_large_finished_and_unfinished_mix():
    """
    Large batch, some batches fully finished, some need to add unfinished.
    """
    batch_size = 10
    beam_size = 8
    eot = 12
    decoder = BeamSearchDecoder(beam_size=beam_size, eot=eot, inference=Inference())
    finished_sequences = []
    for i in range(batch_size):
        if i % 2 == 0:
            # Even: all finished
            seqs = [([i,j], float(j)) for j in range(beam_size)]
        else:
            # Odd: only 2 finished
            seqs = [([i,0], 0.1), ([i,1], 0.2)]
        finished_sequences.append(seqs)
    decoder.finished_sequences = make_finished_sequences(batch_size, finished_sequences)
    # 10 unfinished per batch
    preceding_tokens = torch.arange(batch_size*10*2).reshape(batch_size, 10, 2)
    sum_logprobs = torch.arange(batch_size*10).reshape(batch_size, 10).float()
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)
    for i in range(batch_size):
        if i % 2 == 0:
            # All finished, should not include any with eot
            for t in tokens[i]:
                pass
        else:
            # Should include the two finished, and six unfinished with eot
            seqs = [tuple(t.tolist()) for t in tokens[i]]
            # The rest should end with eot
            unfinished_count = sum(1 for s in seqs if s[-1] == eot)

def test_finalize_large_eot_value():
    """
    Large eot value, to ensure no overflow or type issues.
    """
    beam_size = 2
    eot = 2**31 - 1  # max 32-bit signed int
    decoder = BeamSearchDecoder(beam_size=beam_size, eot=eot, inference=Inference())
    finished_sequences = [
        [([1,2], 0.1)]
    ]
    decoder.finished_sequences = make_finished_sequences(1, finished_sequences)
    preceding_tokens = torch.tensor([[[3,4],[5,6]]])
    sum_logprobs = torch.tensor([[0.9, 0.8]])
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)
    seqs = [tuple(t.tolist()) for t in tokens[0]]

def test_finalize_large_token_sequences():
    """
    Large token sequences (long sequences).
    """
    beam_size = 3
    eot = 0
    decoder = BeamSearchDecoder(beam_size=beam_size, eot=eot, inference=Inference())
    # One finished, two unfinished with long sequences
    long_seq = list(range(100))
    finished_sequences = [
        [(long_seq, 1.0)]
    ]
    decoder.finished_sequences = make_finished_sequences(1, finished_sequences)
    preceding_tokens = torch.stack([torch.arange(100,200), torch.arange(200,300)]).reshape(1,2,100)
    sum_logprobs = torch.tensor([[2.0, 3.0]])
    tokens, logprobs = decoder.finalize(preceding_tokens, sum_logprobs)
    # All sequences should be length 100 or 101
    for t in tokens[0]:
        pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-BeamSearchDecoder.finalize-mayry8p6 and push.