fix(bridge): decode every batch row in generate_stream

tests/integration/model_bridge/test_generate_stream_batched.py

@@ -0,0 +1,48 @@
+"""Batched ``generate_stream(return_type="str")`` decodes every row of the batch.
+
+The streaming decoder used to decode ``tokens[0]`` only, so a batched call
+yielded the first sequence's text for the whole batch and the other rows were
+never observable. See ``generate()``, which decodes each row and unwraps a
+single-row batch — ``generate_stream`` now matches that contract.
+"""
+
+PROMPTS = ["The capital of France is", "My favourite colour is"]
+
+
+def test_batched_stream_yields_one_string_per_row(distilgpt2_bridge):
+    chunks = list(
+        distilgpt2_bridge.generate_stream(
+            PROMPTS,
+            max_new_tokens=4,
+            max_tokens_per_yield=2,
+            do_sample=False,
+            verbose=False,
+            return_type="str",
+        )
+    )
+
+    assert chunks, "expected at least one yield"
+    for chunk in chunks:
+        assert isinstance(chunk, list)
+        assert len(chunk) == len(PROMPTS)
+        assert all(isinstance(text, str) for text in chunk)
+
+    # The first yield carries the input tokens, so each row must echo its own prompt.
+    for prompt, streamed in zip(PROMPTS, chunks[0]):
+        assert prompt in streamed
+
+
+def test_single_prompt_stream_stays_a_bare_string(distilgpt2_bridge):
+    """A one-element batch keeps the scalar contract shared with generate()."""
+    chunks = list(
+        distilgpt2_bridge.generate_stream(
+            [PROMPTS[0]],
+            max_new_tokens=2,
+            do_sample=False,
+            verbose=False,
+            return_type="str",
+        )
+    )
+
+    assert chunks
+    assert all(isinstance(chunk, str) for chunk in chunks)

tests/unit/model_bridge/test_generate_stream_batch_decode.py

@@ -0,0 +1,105 @@
+"""``generate_stream(return_type="str")`` decodes every row of the batch.
+
+The streaming decoder used to call ``tokenizer.decode(tokens[0])``, so a batched
+stream repeated the first sequence's text for the whole batch. These tests drive
+the yield/decode bookkeeping with a stub token stream — no model load — and pin
+both halves of the contract: a one-row batch still yields a bare string, and a
+larger batch yields one string per row, as ``generate()`` does.
+"""
+
+from types import MethodType, SimpleNamespace
+
+import torch
+
+from transformer_lens.model_bridge.bridge import TransformerBridge
+
+
+class _Tokenizer:
+    eos_token_id = None
+    pad_token_id = 0
+    padding_side = "right"
+
+    def __call__(self, inputs, **kwargs):
+        if isinstance(inputs, str) or len(inputs) == 1:
+            return {"input_ids": torch.tensor([[11, 12]])}
+        return {"input_ids": torch.tensor([[11, 12], [21, 22]])}
+
+    def decode(self, tokens, **kwargs):
+        return " ".join(str(token) for token in tokens.tolist())
+
+
+class _Model(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.config = SimpleNamespace(is_encoder_decoder=False)
+
+
+def _make_bridge(generated_steps: list[list[int]]) -> TransformerBridge:
+    bridge = object.__new__(TransformerBridge)
+    torch.nn.Module.__init__(bridge)
+    bridge.cfg = SimpleNamespace(device=torch.device("cpu"), eos_token_id=None)
+    bridge.tokenizer = _Tokenizer()
+    bridge.__dict__["original_model"] = _Model()
+    bridge._ensure_generation_supported = MethodType(lambda self, api: None, bridge)
+    bridge._resolve_generation_caching = MethodType(lambda self, requested, batched: False, bridge)
+
+    def fake_generate_tokens(self, *args, **kwargs):
+        for step, generated_tokens in enumerate(generated_steps):
+            yield torch.tensor(generated_tokens), None, step == len(generated_steps) - 1
+
+    bridge._generate_tokens = MethodType(fake_generate_tokens, bridge)
+    return bridge
+
+
+def test_batched_string_stream_decodes_every_row() -> None:
+    chunks = list(
+        _make_bridge([[13, 23], [14, 24]]).generate_stream(
+            ["first", "second"],
+            max_new_tokens=2,
+            max_tokens_per_yield=1,
+            stop_at_eos=False,
+            do_sample=False,
+            use_past_kv_cache=False,
+            return_type="input",
+            verbose=False,
+        )
+    )
+
+    assert chunks == [
+        ["11 12 13", "21 22 23"],
+        ["14", "24"],
+    ]
+
+
+def test_batched_token_stream_decodes_every_row() -> None:
+    chunks = list(
+        _make_bridge([[33, 43]]).generate_stream(
+            torch.tensor([[31, 32], [41, 42]]),
+            max_new_tokens=1,
+            max_tokens_per_yield=99,
+            stop_at_eos=False,
+            do_sample=False,
+            use_past_kv_cache=False,
+            return_type="str",
+            verbose=False,
+        )
+    )
+
+    assert chunks == [["31 32 33", "41 42 43"]]
+
+
+def test_single_string_stream_remains_scalar() -> None:
+    chunks = list(
+        _make_bridge([[13]]).generate_stream(
+            "first",
+            max_new_tokens=1,
+            max_tokens_per_yield=99,
+            stop_at_eos=False,
+            do_sample=False,
+            use_past_kv_cache=False,
+            return_type="input",
+            verbose=False,
+        )
+    )
+
+    assert chunks == ["11 12 13"]

transformer_lens/model_bridge/bridge.py

@@ -4374,7 +4374,7 @@ def generate_stream(
         verbose: bool = True,
         stop_strings: Optional[Union[str, List[str]]] = None,
         stopping_criteria: Optional[Any] = None,
-    ) -> Generator[Union[torch.Tensor, str], None, None]:
+    ) -> Generator[Union[torch.Tensor, str, List[str]], None, None]:
         """Stream tokens from the model as they are generated.
 
         Yields batches of tokens progressively during generation rather than
@@ -4410,9 +4410,11 @@ def generate_stream(
                 (scores is the step's logits). See generate() for the full contract.
 
         Yields:
-            Token tensors [batch, seq_len] or strings, accumulated up to
-            max_tokens_per_yield tokens between yields. First yield includes
-            the input tokens; subsequent yields contain only new tokens.
+            Token tensors [batch, seq_len], or decoded text when return_type='str' -
+            a bare string for a single sequence and one string per batch row for a
+            larger batch, matching generate(). Chunks accumulate up to
+            max_tokens_per_yield tokens between yields; the first yield includes the
+            input tokens and subsequent yields contain only new tokens.
         """
         self._ensure_generation_supported("generate_stream")
         # --- Input parsing (mirrors generate()) ---
@@ -4529,10 +4531,13 @@ def generate_stream(
 
         def _maybe_decode(
             tokens: torch.Tensor,
-        ) -> Union[torch.Tensor, str]:
+        ) -> Union[torch.Tensor, str, List[str]]:
             if return_type == "str":
                 assert self.tokenizer is not None
-                return self.tokenizer.decode(tokens[0], skip_special_tokens=True)
+                decoded_texts = [
+                    self.tokenizer.decode(row, skip_special_tokens=True) for row in tokens
+                ]
+                return decoded_texts[0] if len(decoded_texts) == 1 else decoded_texts
             return tokens
 
         try: