From f6947490900e158e41cb77abc5f92509d054944b Mon Sep 17 00:00:00 2001 From: ashutosh0x Date: Mon, 5 Jan 2026 03:14:45 +0530 Subject: [PATCH] Fix: Handle empty prompt arrays in make_seq2seq_fields (#483) --- gemma/gm/data/_functional.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/gemma/gm/data/_functional.py b/gemma/gm/data/_functional.py index b9e4fef5..480e5657 100644 --- a/gemma/gm/data/_functional.py +++ b/gemma/gm/data/_functional.py @@ -138,9 +138,18 @@ def make_seq2seq_fields( sequence = np.concatenate([prompt, response]) # Create the loss mask. + if len(prompt) > 0: + prompt_mask = np.zeros((len(prompt) - 1,), dtype=np.bool_) + response_mask = np.ones((len(response),), dtype=np.bool_) + else: + # If no prompt, all targets come from response (shifted by 1) + prompt_mask = np.zeros((0,), dtype=np.bool_) + # Ensure non-negative shape if response is also empty or length 1 + response_mask = np.ones((max(0, len(response) - 1),), dtype=np.bool_) + target_mask = np.concatenate([ - np.zeros((len(prompt) - 1,), dtype=np.bool_), - np.ones((len(response),), dtype=np.bool_), + prompt_mask, + response_mask, ]) return Seq2SeqFields(