fix generate

model.py

@@ -2130,6 +2130,132 @@ class CustomWhisperGenerationMixin(WhisperGenerationMixin):
                         for i in range(len(outputs.encoder_hidden_states))
                     )
             return outputs
+        
+        def _pad_to_max_length(
+            current_segments,
+            pad_token_id,
+            device,
+            padding_side="right",
+            padding="longest",
+            bos_token_tensor=None,
+            cut_off_length=None,
+            return_token_timestamps=False,
+            force_unique_generate_call=False,
+        ):
+            max_total_length = 0
+            sequences = []
+            token_timestamps_list = []
+
+            if padding_side not in ["right", "left"]:
+                raise ValueError(f"`padding_side` must be either 'right' or 'left', not {padding_side}")
+
+            if padding not in ["longest", "max_length"]:
+                raise ValueError(f"`padding` must be either 'longest' or 'max_length', not {padding}")
+            elif padding == "max_length" and cut_off_length is None:
+                raise ValueError("`cut_off_length` must be specified when `padding='max_length'`")
+
+            if force_unique_generate_call:
+                sequences_list = []
+                timestamps_list = []
+                for segments in current_segments:
+                    result = segments[0]["result"]
+                    sequences_list.append(result if isinstance(result, torch.Tensor) else result["sequences"])
+                    if return_token_timestamps:
+                        timestamps_list.append(result["token_timestamps"])
+
+                sequences = torch.stack(sequences_list, dim=0)
+                if return_token_timestamps:
+                    token_timestamps = torch.stack(timestamps_list, dim=0)
+                    return sequences, token_timestamps
+                return sequences
+
+            for current_segment_list in current_segments:
+                if current_segment_list is not None and len([d["tokens"] for d in current_segment_list]) > 0:
+                    sequence = torch.cat([d["tokens"] for d in current_segment_list], dim=-1)
+                    if return_token_timestamps:
+                        token_timestamps = torch.cat(
+                            [d["result"]["token_timestamps"][d["idxs"][0] : d["idxs"][1]] for d in current_segment_list],
+                            dim=-1,
+                        )
+
+                    if cut_off_length is not None:
+                        sequence = sequence[-cut_off_length:]
+                        if return_token_timestamps:
+                            token_timestamps = token_timestamps[-cut_off_length:]
+
+                    if bos_token_tensor is not None:
+                        sequence = torch.cat([bos_token_tensor, sequence])
+                        if return_token_timestamps:
+                            token_timestamps = torch.cat(
+                                [torch.ones_like(bos_token_tensor, device=device) * 0.0, token_timestamps]
+                            )
+                    sequences.append(sequence)
+                    if return_token_timestamps:
+                        token_timestamps_list.append(token_timestamps)
+                    max_total_length = max(max_total_length, len(sequences[-1]))
+                elif bos_token_tensor is not None:
+                    sequences.append(bos_token_tensor)
+                    if return_token_timestamps:
+                        token_timestamps_list.append(torch.ones_like(bos_token_tensor, device=device) * 0.0)
+                else:
+                    sequences.append(torch.tensor([], device=device))
+                    if return_token_timestamps:
+                        token_timestamps_list.append(torch.tensor([], device=device))
+
+            max_total_length = cut_off_length + 1 if padding == "max_length" else max_total_length
+            for i in range(len(current_segments)):
+                pad_length = max_total_length - len(sequences[i])
+                pad = (0, pad_length) if padding_side == "right" else (pad_length, 0)
+
+                sequences[i] = F.pad(sequences[i], pad=pad, value=pad_token_id)
+                if return_token_timestamps:
+                    token_timestamps_list[i] = F.pad(
+                        token_timestamps_list[i],
+                        pad=pad,
+                        value=token_timestamps_list[i][-1] if len(token_timestamps_list[i]) > 0 else 0.0,
+                    )
+
+            sequences = torch.stack(sequences, dim=0)
+
+            if return_token_timestamps:
+                token_timestamps = torch.stack(token_timestamps_list, dim=0)
+                return sequences, token_timestamps
+            else:
+                return sequences
+
+        padded_outputs = _pad_to_max_length(
+            current_segments=final_segments,
+            pad_token_id=generation_config.pad_token_id,
+            device=self.device,
+            padding_side="right",
+            return_token_timestamps=return_token_timestamps,
+            force_unique_generate_call=force_unique_generate_call,
+        )
+
+        if return_dict_in_generate and generation_config.return_dict_in_generate:
+            logger.warning_once(
+                "You have passed `return_dict_in_generate=True` and `return_timestamps=True`, this automatically sets `return_segments=True` to access the resuls of the underlying calls to GenerationMixin's generate in the returned `segments`."
+            )
+            return_segments = True
+        elif not return_segments and not return_token_timestamps:
+            return padded_outputs
+
+        if return_token_timestamps:
+            sequences, token_timestamps = padded_outputs
+            outputs = {
+                "sequences": sequences,
+                "token_timestamps": token_timestamps,
+            }
+        else:
+            sequences = padded_outputs
+            outputs = {
+                "sequences": sequences,
+            }
+
+        if return_segments:
+            outputs["segments"] = final_segments
+
+        return outputs
 
 @add_start_docstrings(
     "The Whisper Model with a language modeling head. Can be used for automatic speech recognition.",