Text-generation pipeline crashes with batched generation metadata and silently drops beam scores #50209

System Info

• Transformers: 5.23.0.dev0 (development version installed from GitHub main)
• PyTorch: 2.11.0+cu128
• Python: 3.13
• Platform: Linux
• GPU: NVIDIA Tesla T4
• Models tested: openai-community/gpt2, facebook/opt-125m
• Precision: FP32

Reproduction

While testing batched text generation with scoring metadata noticed two related problems in TextGenerationPipeline.

First, requesting structured generation outputs with return_dict_in_generate=True and output_scores=True causes batched pipeline inference to crash with KeyError: 0, even though direct model.generate() works.

Second, requesting multiple beam candidates silently drops sequences_scores from the pipeline results, although the underlying model returns them correctly.

Both problems reproduce on Transformers 5.23.0.dev0 with GPT-2 and OPT-125M.

Minimal reproduction

import torch
from transformers import pipeline

p = pipeline(
    "text-generation",
    model="openai-community/gpt2",
    device=0 if torch.cuda.is_available() else -1,
)
p.tokenizer.pad_token = p.tokenizer.eos_token
p.tokenizer.padding_side = "left"

prompts = [
    "Artificial intelligence is",
    "The future of science is",
]

kwargs = dict(
    max_new_tokens=4,
    do_sample=False,
    num_beams=3,
    num_return_sequences=1,
    return_dict_in_generate=True,
    output_scores=True,
)

# Control: direct generation succeeds
inputs = p.tokenizer(
    prompts, padding=True, return_tensors="pt"
).to(p.model.device)

with torch.inference_mode():
    direct = p.model.generate(**inputs, **kwargs)

print(direct.sequences_scores.shape)
# torch.Size([2])

# Bug 1: pipeline fails with actual batching
p(
    prompts,
    batch_size=2,
    return_full_text=False,
    **kwargs,
)
# KeyError: 0

The failure occurs in:

transformers/pipelines/pt_utils.py

PipelineIterator.loader_batch_item()

elif isinstance(
    element[self._loader_batch_index], torch.Tensor
):
KeyError: 0

The second issue can be reproduced independently:

kwargs["num_return_sequences"] = 2

inputs = p.tokenizer(
    [prompts[0]], return_tensors="pt"
).to(p.model.device)

with torch.inference_mode():
    direct = p.model.generate(**inputs, **kwargs)

print(direct.sequences_scores.shape)
# torch.Size([2])

out = p(
    [prompts[0]],
    batch_size=1,
    return_full_text=False,
    **kwargs,
)

print([list(record.keys()) for record in out[0]])
# [['generated_text'], ['generated_text']]

Both candidates are returned, but their sequence scores disappear without an error.

Cross-model results

I tested 12 configurations across GPT-2 and OPT-125M.

Summary:

• 12 configurations tested
• 6 crashes
• 2 silent sequence-score losses
• 4 passing controls
• Direct model.generate() succeeded in every configuration

Suspected root causes

1. Nested generation metadata is not properly unbatched

In src/transformers/pipelines/text_generation.py, _forward() stores extra generation outputs inside an additional_outputs dictionary.

During batched pipeline execution, PipelineIterator.loader_batch_item() in src/transformers/pipelines/pt_utils.py attempts integer indexing on this nested dictionary. This produces KeyError: 0.

2. Sequence scores are not correctly aligned with candidates

TextGenerationPipeline._forward() reshapes generation-level tensors to include batch and candidate dimensions.

However, postprocess() only extracts metadata tensors when their first dimension matches the number of candidates.

This causes sequences_scores to be omitted when multiple sequences are returned, because the tensor still contains an additional batch dimension.

Existing regression coverage

The current test_return_dict_in_generate in tests/pipelines/test_pipelines_text_generation.py covers structured generation outputs, but does not explicitly test batch_size > 1 or multiple returned beam candidates.

Adding those cases should help prevent regressions.

The two failures appear related to how structured generation outputs are propagated through the pipeline rather than to the underlying model architectures.

Expected behavior

The text-generation pipeline should correctly handle structured generation outputs for both single-input and batched inference.

Specifically:

1. Calling pipeline() with return_dict_in_generate=True, output_scores=True, and batch_size > 1 should complete without raising KeyError: 0.
2. When beam search returns multiple candidates, each generated sequence should retain its corresponding sequences_scores value.
3. Generation metadata should remain correctly associated with each input and candidate after batching and postprocessing.
4. Pipeline results should preserve applicable generation metadata without silently dropping supported fields.
5. Ordinary generation behavior without structured metadata should remain unchanged.

Direct model.generate() already returns the expected sequence scores in the reproduced configurations, so the pipeline should preserve that information while converting the outputs into its standard result format.