update patch for generation method

model.py

@@ -1801,25 +1801,13 @@ class WhisperModel(WhisperPreTrainedModel):
             encoder_attentions=encoder_outputs.attentions,
         )
 
+        
 
-@add_start_docstrings(
-    "The Whisper Model with a language modeling head. Can be used for automatic speech recognition.",
-    WHISPER_START_DOCSTRING,
-)
-class CustomWhisperForConditionalGeneration(WhisperGenerationMixin, WhisperPreTrainedModel):
-    base_model_prefix = "model"
-    _tied_weights_keys = ["proj_out.weight"]
-
-    def __init__(self, config: CustomWhisperConfig):
-        super().__init__(config)
-        self.model = WhisperModel(config)
-        self.proj_out = nn.Linear(config.d_model, config.vocab_size, bias=False)
-        self.max_target_positions = config.max_target_positions
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    # CUSTOM (patch the generation method)       
+# CUSTOM (patch the generation method)    
+class CustomWhisperGenerationMixin(WhisperGenerationMixin):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+           
     def generate(
         self,
         input_features: Optional[torch.Tensor] = None,
@@ -2142,6 +2130,23 @@ class CustomWhisperForConditionalGeneration(WhisperGenerationMixin, WhisperPreTr
                         for i in range(len(outputs.encoder_hidden_states))
                     )
             return outputs
+
+@add_start_docstrings(
+    "The Whisper Model with a language modeling head. Can be used for automatic speech recognition.",
+    WHISPER_START_DOCSTRING,
+)
+class CustomWhisperForConditionalGeneration(CustomWhisperGenerationMixin, WhisperPreTrainedModel):
+    base_model_prefix = "model"
+    _tied_weights_keys = ["proj_out.weight"]
+
+    def __init__(self, config: CustomWhisperConfig):
+        super().__init__(config)
+        self.model = WhisperModel(config)
+        self.proj_out = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.max_target_positions = config.max_target_positions
+
+        # Initialize weights and apply final processing
+        self.post_init()
         
         def _pad_to_max_length(
             current_segments,