System Info

• transformers version: 4.39.0 (also verified present in 4.53.3 and main branch)
• Python version: 3.10
• NumPy version: 1.x
• OS: Linux (AWS SageMaker)

Who can help?

@Narsil

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder
• My own task or dataset

Reproduction

The decode_spans function in src/transformers/pipelines/question_answering.py has an off-by-one boundary check on line 81:

scores_flat = candidates.flatten()
if topk == 1:
    idx_sort = [np.argmax(scores_flat)]
elif len(scores_flat) < topk:          # <-- should be <=
    idx_sort = np.argsort(-scores_flat)
else:
    idx = np.argpartition(-scores_flat, topk)[0:topk]  # crashes when len == topk
    idx_sort = idx[np.argsort(-scores_flat[idx])]

np.argpartition(array, kth) requires 0 <= kth < len(array). When len(scores_flat) == topk, the guard len(scores_flat) < topk evaluates to False, and argpartition is called with kth == len(array), which is out of bounds.

Minimal reproduction:

import numpy as np

def decode_spans_bug(seq_len=10, topk=100, max_answer_len=15):
    """Reproduces the bug when seq_len² == topk"""
    start = np.random.rand(1, seq_len)
    end = np.random.rand(1, seq_len)
    outer = np.matmul(np.expand_dims(start, -1), np.expand_dims(end, 1))
    candidates = np.tril(np.triu(outer), max_answer_len - 1)
    scores_flat = candidates.flatten()
    
    print(f"len(scores_flat)={len(scores_flat)}, topk={topk}")
    print(f"len < topk: {len(scores_flat) < topk}")
    
    # This is what decode_spans does:
    idx = np.argpartition(-scores_flat, topk)[0:topk]  # CRASH

decode_spans_bug()  # ValueError: kth(=100) out of bounds (100)

Since scores_flat has seq_len² elements (from the outer product matrix), this triggers when topk is a perfect square and seq_len = sqrt(topk). For example:

• top_k=100 crashes when seq_len=10
• top_k=400 crashes when seq_len=20

This occurs in practice when using top_k > 1 with very short input contexts (e.g., a short question with minimal context text).

Expected behavior

The function should handle the boundary case where len(scores_flat) == topk without crashing. The fix is changing < to <= on line 81:

elif len(scores_flat) <= topk:    # fixed: catches the boundary case
    idx_sort = np.argsort(-scores_flat)

When len(scores_flat) == topk, we want ALL elements anyway, so argsort is both correct and sufficient — no need for argpartition.