fix: add compatibility with Transformers 4.55.0

SCBs.py

@@ -2,7 +2,18 @@ import torch
 from torch import nn
 from transformers import WhisperConfig
 from transformers.activations import ACT2FN
-from transformers.models.whisper.modeling_whisper import WHISPER_ATTENTION_CLASSES
+# Compatibility fallback for Transformers versions
+# - Transformers <= 4.38: WHISPER_ATTENTION_CLASSES is available
+# - Transformers >= 4.39 (including 4.55.0): WHISPER_ATTENTION_CLASSES removed, use WhisperAttention dispatcher
+try:
+    from transformers.models.whisper.modeling_whisper import WHISPER_ATTENTION_CLASSES
+except ImportError:
+    from transformers.models.whisper.modeling_whisper import WhisperAttention
+    WHISPER_ATTENTION_CLASSES = {
+        "eager": WhisperAttention,
+        "sdpa": WhisperAttention,
+        "flash_attention_2": WhisperAttention,
+    }
 import torch.nn.functional as F
 from .coattention import CoAttention
 from .layers import CustomLinear, CustomDiagonalLinear, Gate

encoder.py

@@ -1,7 +1,19 @@
 import torch
 from torch import nn
 from transformers.modeling_outputs import CausalLMOutput, BaseModelOutput
-from transformers.models.whisper.modeling_whisper import WhisperEncoder, WhisperEncoderLayer, WHISPER_ATTENTION_CLASSES
+# Compatibility fallback for Transformers versions
+# - Transformers <= 4.38: WHISPER_ATTENTION_CLASSES is available
+# - Transformers >= 4.39 (including 4.55.0): WHISPER_ATTENTION_CLASSES removed, use WhisperAttention dispatcher
+from transformers.models.whisper.modeling_whisper import WhisperEncoder, WhisperEncoderLayer
+try:
+    from transformers.models.whisper.modeling_whisper import WHISPER_ATTENTION_CLASSES
+except ImportError:
+    from transformers.models.whisper.modeling_whisper import WhisperAttention
+    WHISPER_ATTENTION_CLASSES = {
+        "eager": WhisperAttention,
+        "sdpa": WhisperAttention,
+        "flash_attention_2": WhisperAttention,
+    }
 
 from .FDDT import FDDT
 from .config import DiCoWConfig
@@ -244,7 +256,7 @@ class DiCoWEncoder(WhisperEncoder):
                 layer_head_mask=None,
             )
         elif hasattr(self, "additional_self_attention_layer"):
-            inter_output, _, __ = self.additional_self_attention_layer(
+            inter_output, _ = self.additional_self_attention_layer(
                 outputs.last_hidden_state,
                 attention_mask=None,
                 output_attentions=output_attentions,

generation.py

@@ -25,9 +25,10 @@ from transformers.generation.stopping_criteria import (
     StoppingCriteriaList,
 )
 from transformers.generation.utils import GenerateBeamOutput, BeamScorer, GenerateBeamDecoderOnlyOutput, \
-    stack_model_outputs, GenerateBeamEncoderDecoderOutput, _split_model_inputs, GenerateNonBeamOutput, \
+    GenerateBeamEncoderDecoderOutput, GenerateNonBeamOutput, \
     GenerateEncoderDecoderOutput, GenerateDecoderOnlyOutput
 from transformers.modeling_outputs import BaseModelOutput
+from transformers.utils import ModelOutput
 from transformers.models.whisper.modeling_whisper import (
     WhisperForConditionalGeneration,
 )
@@ -42,6 +43,93 @@ logging.set_verbosity_debug()
 logger = logging.get_logger("transformers")
 
 
+# Backport of transformers 4.40 helpers removed in 4.55
+def _split(data, full_batch_size: int, split_size: int = None):
+    if data is None:
+        return [None] * (full_batch_size // split_size)
+    if isinstance(data, torch.Tensor):
+        return [data[i: i + split_size] for i in range(0, full_batch_size, split_size)]
+    elif isinstance(data, tuple):
+        if isinstance(data[0], tuple):
+            return [
+                tuple(tuple(tensor[i: i + split_size] for tensor in inner_tuple) for inner_tuple in data)
+                for i in range(0, full_batch_size, split_size)
+            ]
+        else:
+            return [
+                tuple(sub_tensor[i: i + split_size] for sub_tensor in data)
+                for i in range(0, full_batch_size, split_size)
+            ]
+    else:
+        raise ValueError(f"Unexpected attribute type: {type(data)}")
+
+
+def _split_model_inputs(
+    model_input: Union[ModelOutput, Dict], split_size: int, full_batch_size: int
+) -> List[Union[ModelOutput, Dict]]:
+    if model_input is None:
+        return [model_input] * (full_batch_size // split_size)
+    model_output_cls = type(model_input)
+    if (full_batch_size % split_size) != 0:
+        raise ValueError("`full_batch_size` must be divisible by `split_size`")
+    if split_size > full_batch_size:
+        raise ValueError("`split_size` must be smaller or equal to `full_batch_size`")
+    keys = (
+        model_input.__dataclass_fields__.keys() if hasattr(model_input, "__dataclass_fields__") else model_input.keys()
+    )
+    keys = [k for k in keys if k in model_input]
+    bool_keys = [k for k in keys if isinstance(model_input[k], bool) or k == "cache_position"]
+    keys_to_ignore = ["cache_position", "encoder_outputs", "num_logits_to_keep"]
+    non_bool_keys = [k for k in keys if not isinstance(model_input[k], bool) and k not in keys_to_ignore]
+    data_split_list = [
+        {k: _split(model_input[k], full_batch_size, split_size)[i] for k in non_bool_keys}
+        for i in range(full_batch_size // split_size)
+    ]
+    bool_data = {k: model_input[k] for k in bool_keys}
+    if "encoder_outputs" in model_input:
+        encoder_outputs_split = _split_model_inputs(model_input["encoder_outputs"], split_size, full_batch_size)
+        data_split_list = [
+            {**data_split, "encoder_outputs": encoder_outputs_split[i]} for i, data_split in enumerate(data_split_list)
+        ]
+    if "num_logits_to_keep" in model_input:
+        data_split_list = [
+            {**data_split, "num_logits_to_keep": model_input["num_logits_to_keep"]} for data_split in data_split_list
+        ]
+    return [model_output_cls(**data_split, **bool_data) for data_split in data_split_list]
+
+
+def stack_model_outputs(model_outputs: List[ModelOutput]) -> ModelOutput:
+    if not model_outputs:
+        raise ValueError("Input list is empty.")
+    model_output_cls = type(model_outputs[0])
+    if not all(isinstance(obj, model_output_cls) for obj in model_outputs):
+        raise ValueError("All elements in the list should be of the same type.")
+
+    def _concat(data):
+        if any(d is None for d in data):
+            return None
+        if isinstance(data[0], torch.Tensor):
+            return torch.cat(data, dim=0)
+        elif isinstance(data[0], tuple):
+            if isinstance(data[0][0], tuple):
+                return tuple(
+                    tuple(torch.cat([attr[i][j] for attr in data], dim=0) for j in range(len(data[0][0])))
+                    for i in range(len(data[0]))
+                )
+            else:
+                return tuple(torch.cat([attr[i] for attr in data], dim=0) for i in range(len(data[0])))
+        elif isinstance(data[0], (int, float)):
+            return torch.tensor(data)
+        else:
+            raise ValueError(f"Unexpected attribute type: {type(data[0])}")
+
+    concatenated_data = {
+        k: _concat([getattr(model_output, k) for model_output in model_outputs])
+        for k in model_output_cls.__dataclass_fields__.keys()
+    }
+    return model_output_cls(**concatenated_data)
+
+
 class DiCoWGenerationMixin(WhisperForConditionalGeneration):
     def _prepare_encoder_decoder_kwargs_for_generation(
             self, inputs_tensor: torch.Tensor, model_kwargs, model_input_name, generation_config,
@@ -55,6 +143,8 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
         model_kwargs = super()._prepare_encoder_decoder_kwargs_for_generation(
             inputs_tensor, model_kwargs, model_input_name, generation_config
         )
+        # Ensure output_hidden_states is in model_kwargs
+        model_kwargs["output_hidden_states"] = True
         self.encoder_logits = model_kwargs["encoder_outputs"].logits
 
         return model_kwargs
@@ -153,6 +243,9 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
             input_features=input_features, input_stride=input_stride, kwargs=kwargs
         )
         is_shortform = total_input_frames <= num_segment_frames
+        # Store for use in generate_with_fallback when called via parent's generate (shortform path)
+        self._num_segment_frames = num_segment_frames
+        self._max_frames = torch.ones((batch_size,), dtype=torch.long) * total_input_frames
 
         if is_shortform:
             # warn user of ignored inputs
@@ -170,7 +263,6 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
         self._set_return_outputs(
             return_dict_in_generate=return_dict_in_generate,
             return_token_timestamps=return_token_timestamps,
-            is_shortform=is_shortform,
             logprob_threshold=logprob_threshold,
             generation_config=generation_config,
         )
@@ -181,7 +273,8 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
             language=language, task=task, is_multilingual=is_multilingual, generation_config=generation_config
         )
         self._set_num_frames(
-            return_token_timestamps=return_token_timestamps, generation_config=generation_config, kwargs=kwargs
+            return_token_timestamps=return_token_timestamps, generation_config=generation_config,
+            attention_mask=attention_mask, kwargs=kwargs
         )
         self._set_thresholds_and_condition(
             generation_config=generation_config,
@@ -278,8 +371,10 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
         batch_size = input_features.shape[0]
 
         max_frames, seek = self._retrieve_max_frames_and_seek(
-            batch_size=batch_size, attention_mask=attention_mask, total_input_frames=total_input_frames
+            batch_size=batch_size, attention_mask=attention_mask, total_input_frames=total_input_frames,
+            is_shortform=is_shortform,
         )
+        self._max_frames = max_frames
 
         # 6.2 Preppare running variables, list for generation
         cur_bsz = batch_size
@@ -349,7 +444,7 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
                     proc.set_begin_index(decoder_input_ids.shape[-1])
 
             # 6.8 Run generate with fallback
-            seek_sequences, seek_outputs, should_skip, do_condition_on_prev_tokens = self.generate_with_fallback(
+            seek_sequences, seek_outputs, should_skip, do_condition_on_prev_tokens, _model_output_type = self.generate_with_fallback(
                 segment_input=segment_input,
                 decoder_input_ids=decoder_input_ids,
                 cur_bsz=cur_bsz,
@@ -732,7 +827,7 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
         num_beams = beam_scorer.num_beams
 
         batch_beam_size, cur_len = input_ids.shape
-        model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs)
+        model_kwargs = self._get_initial_cache_position(input_ids.shape[-1], input_ids.device, model_kwargs)
 
         if num_beams * batch_size != batch_beam_size:
             raise ValueError(
@@ -795,7 +890,8 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
                 )
                 outputs_per_sub_batch = [
                     self(
-                        **inputs_per_sub_batch,
+                        **{k: v for k, v in inputs_per_sub_batch.items()
+                           if k not in ("output_attentions", "output_hidden_states")},
                         return_dict=True,
                         output_attentions=output_attentions,
                         output_hidden_states=output_hidden_states,
@@ -806,8 +902,10 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
                 outputs = stack_model_outputs(outputs_per_sub_batch)
 
             else:  # Unchanged original behavior
+                _beam_model_inputs = {k: v for k, v in model_inputs.items()
+                                      if k not in ("output_attentions", "output_hidden_states")}
                 outputs = self(
-                    **model_inputs,
+                    **_beam_model_inputs,
                     return_dict=True,
                     output_attentions=output_attentions,
                     output_hidden_states=output_hidden_states,
@@ -1034,13 +1132,16 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
         batch_size = input_ids.shape[0]
         this_peer_finished = False
         unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
-        model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs)
+        model_kwargs = self._get_initial_cache_position(input_ids.shape[-1], input_ids.device, model_kwargs)
 
         while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
             # prepare model inputs
             model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
 
             # forward pass to get next token
+            # Pop keys that we pass explicitly to avoid duplicate keyword argument error
+            model_inputs.pop("output_attentions", None)
+            model_inputs.pop("output_hidden_states", None)
             outputs = self(
                 **model_inputs,
                 return_dict=True,
@@ -1186,10 +1287,8 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
             segment_input,
             decoder_input_ids,
             cur_bsz,
-            batch_idx_map,
             seek,
-            num_segment_frames,
-            max_frames,
+            batch_idx_map,
             temperatures,
             generation_config,
             logits_processor,
@@ -1198,36 +1297,46 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
             synced_gpus,
             return_token_timestamps,
             do_condition_on_prev_tokens,
-            kwargs,
+            is_shortform=False,
+            batch_size=None,
+            attention_mask=None,
+            kwargs=None,
+            # Legacy args kept for DiCoW's own generate() call
+            num_segment_frames=None,
+            max_frames=None,
     ):
+        if kwargs is None:
+            kwargs = {}
         kwargs = copy.copy(kwargs)
-        kwargs = self.prepare_kwargs_for_generate(segment_input, cur_bsz, batch_idx_map, seek, num_segment_frames,
-                                                  max_frames, kwargs)
-        seek_sequences, seek_outputs, should_skip, do_condition_on_prev_tokens = super().generate_with_fallback(
-            segment_input,
-            decoder_input_ids,
-            cur_bsz,
-            batch_idx_map,
-            seek,
-            num_segment_frames,
-            max_frames,
-            temperatures,
-            generation_config,
-            logits_processor,
-            stopping_criteria,
-            prefix_allowed_tokens_fn,
-            synced_gpus,
-            return_token_timestamps,
-            do_condition_on_prev_tokens,
-            kwargs,
-        )
+        # Use instance-stored values if not provided (e.g. when called from parent's generate)
+        _num_segment_frames = num_segment_frames if num_segment_frames is not None else getattr(self, "_num_segment_frames", None)
+        _max_frames = max_frames if max_frames is not None else getattr(self, "_max_frames", None)
+        if _num_segment_frames is not None and _max_frames is not None:
+            kwargs = self.prepare_kwargs_for_generate(segment_input, cur_bsz, batch_idx_map, seek,
+                                                      _num_segment_frames, _max_frames, kwargs)
+        seek_sequences, seek_outputs, should_skip, do_condition_on_prev_tokens, _model_output_type = \
+            super().generate_with_fallback(
+                segment_input=segment_input,
+                decoder_input_ids=decoder_input_ids,
+                cur_bsz=cur_bsz,
+                seek=seek,
+                batch_idx_map=batch_idx_map,
+                temperatures=temperatures,
+                generation_config=generation_config,
+                logits_processor=logits_processor,
+                stopping_criteria=stopping_criteria,
+                prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
+                synced_gpus=synced_gpus,
+                return_token_timestamps=return_token_timestamps,
+                do_condition_on_prev_tokens=do_condition_on_prev_tokens,
+                is_shortform=is_shortform,
+                batch_size=batch_size if batch_size is not None else cur_bsz,
+                attention_mask=attention_mask,
+                kwargs=kwargs,
+            )
         self.stno_mask_seek = None
 
-        # for i,  seq in enumerate(seek_outputs):
-        #     print(f"Sequence {i}: {self.tokenizer.decode(seq, decode_with_timestamps=True)}")
-        # print("-"*50)
-
-        return seek_sequences, seek_outputs, should_skip, do_condition_on_prev_tokens
+        return seek_sequences, seek_outputs, should_skip, do_condition_on_prev_tokens, _model_output_type
 
     def _retrieve_init_tokens(self, input_features, batch_size, generation_config, config, num_segment_frames, kwargs):
         def replace_or_add(lst: List[int], num: int, itr: Iterator[int]):
@@ -1264,7 +1373,7 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
         task = getattr(generation_config, "task", None)
         language = getattr(generation_config, "language", None)
 
-        forced_decoder_ids = generation_config.forced_decoder_ids
+        forced_decoder_ids = getattr(generation_config, "forced_decoder_ids", None)
         if forced_decoder_ids is not None:
             if language is None and task is None and forced_decoder_ids[0][1] is None:
                 logger.warning_once(
@@ -1289,7 +1398,8 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
             return forced_decoder_ids
 
         # from v4.39 the forced decoder ids are always None in favour of decoder input ids
-        generation_config.forced_decoder_ids = None
+        if hasattr(generation_config, "forced_decoder_ids"):
+            generation_config.forced_decoder_ids = None
 
         is_lang_id_undefined = len(init_tokens) <= 1 or (len(init_tokens) > 1 and init_tokens[1] is None)
 
@@ -1434,17 +1544,17 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
     ) -> LogitsProcessorList:
         # pylint: disable=no-member
         gen_config_copy = copy.deepcopy(generation_config)
-        gen_config_copy.forced_decoder_ids = None
+        if hasattr(gen_config_copy, "forced_decoder_ids"):
+            gen_config_copy.forced_decoder_ids = None
         processors = super()._get_logits_processor(
-            gen_config_copy,
-            input_ids_seq_length,
-            encoder_input_ids,
-            prefix_allowed_tokens_fn,
-            logits_processor,
-            device,
-            model_kwargs,
-            negative_prompt_ids,
-            negative_prompt_attention_mask,
+            generation_config=gen_config_copy,
+            input_ids_seq_length=input_ids_seq_length,
+            encoder_input_ids=encoder_input_ids,
+            prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
+            logits_processor=logits_processor,
+            model_kwargs=model_kwargs,
+            negative_prompt_ids=negative_prompt_ids,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
         )
         if hasattr(generation_config, "ctc_weight") and generation_config.ctc_weight > 0:
             enc_logits = self.encoder_logits
@@ -1469,8 +1579,8 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
             processors.append(self.ctc_rescorer)
         return processors
 
-    def _retrieve_logit_processors(self, generation_config, logits_processor, begin_index, is_shortform, num_beams,
-                                   device):
+    def _retrieve_logit_processors(self, generation_config, logits_processor, begin_index, is_shortform=False,
+                                   num_beams=1, device=None):
         if generation_config.return_timestamps is True:
             timestamp_processor = WhisperTimeStampLogitsProcessorCustom(generation_config, begin_index=begin_index)
             logits_processor = (
@@ -1627,6 +1737,8 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
             prev_idx,
             idx,
             return_token_timestamps,
+            time_precision_features=None,
+            decoder_input_ids=None,
     ):
         # find the predicted "end of segment" predictions of Whisper
         # "end of segment" predictions occur whenever Whisper predicts a timestamp token
@@ -1718,7 +1830,8 @@ class DiCoWGenerationMixin(WhisperForConditionalGeneration):
 
         return segments, segment_offset
 
-    def _postprocess_outputs(self, seek_outputs, decoder_input_ids, return_token_timestamps, generation_config):
+    def _postprocess_outputs(self, seek_outputs, decoder_input_ids, return_token_timestamps, generation_config,
+                             is_shortform=False, seek=None, batch_idx_map=None):
         # remove all previously passed decoder input ids
         if isinstance(seek_outputs, torch.Tensor):
             seek_outputs = seek_outputs[:, decoder_input_ids.shape[-1]:]

modeling_dicow.py

@@ -101,7 +101,7 @@ class DiCoW(WhisperModel):
         decoder_outputs = self.decoder(
             input_ids=decoder_input_ids,
             attention_mask=decoder_attention_mask,
-            encoder_hidden_states=encoder_outputs.hidden_states[-1],
+            encoder_hidden_states=encoder_outputs.hidden_states[-1] if encoder_outputs.hidden_states is not None else encoder_outputs.last_hidden_state,
             head_mask=decoder_head_mask,
             cross_attn_head_mask=cross_attn_head_mask,
             past_key_values=past_key_values,
@@ -122,7 +122,7 @@ class DiCoW(WhisperModel):
             decoder_hidden_states=decoder_outputs.hidden_states,
             decoder_attentions=decoder_outputs.attentions,
             cross_attentions=decoder_outputs.cross_attentions,
-            encoder_last_hidden_state=encoder_outputs.hidden_states[-1],
+            encoder_last_hidden_state=encoder_outputs.hidden_states[-1] if encoder_outputs.hidden_states is not None else encoder_outputs.last_hidden_state,
             encoder_hidden_states=encoder_outputs.hidden_states,
             encoder_attentions=encoder_outputs.attentions,
             encoder_logits=encoder_outputs.logits,
@@ -240,6 +240,7 @@ class DiCoWForConditionalGeneration(DiCoWGenerationMixin, WhisperForConditionalG
             output_hidden_states: Optional[bool] = None,
             return_dict: Optional[bool] = None,
             is_valid: Optional[bool] = None,
+            cache_position: Optional[torch.LongTensor] = None,
     ) -> Union[Tuple[torch.Tensor], Seq2SeqLMOutput]:
         r"""
         labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):

utils.py

@@ -31,9 +31,8 @@ class WhisperTimeStampLogitsProcessorCustom(WhisperTimeStampLogitsProcessor):
             else getattr(generate_config, "_detect_timestamp_from_logprob", True)
         )
 
-        num_forced_ids = (
-            len(generate_config.forced_decoder_ids) if generate_config.forced_decoder_ids is not None else 0
-        )
+        _forced_decoder_ids = getattr(generate_config, "forced_decoder_ids", None)
+        num_forced_ids = len(_forced_decoder_ids) if _forced_decoder_ids is not None else 0
         self.begin_index = begin_index or (num_forced_ids + 1)
 
         self.max_initial_timestamp_index = getattr(generate_config, "max_initial_timestamp_index", None)