Update model.py

model.py

@@ -1064,65 +1064,6 @@ class WhisperEncoder(WhisperPreTrainedModel):
                 return embed_pos[:max_pos_len]
             else:
                 return embed_pos[-max_pos_len:]
-    
-    # CUSTOM (Monkeypatch the generation method)
-    def patch_generate():
-        """
-        Monkey patches the WhisperGenerationMixin to use dynamic stride calculation
-        """
-        original_generate = WhisperGenerationMixin.generate
-        
-        def get_conv_stride(self):
-            """Calculate total stride of all conv layers"""
-            total_stride = 1
-            for layer in self.model.encoder.conv_layers:
-                total_stride *= layer.stride[0]
-            return total_stride
-        
-        def generate_wrapper(self, *args, **kwargs):
-            # Store the original function logic
-            original_code = original_generate.__code__
-            
-            # Create a modified version of the function that uses our stride calculation
-            modified_code = types.CodeType(
-                original_code.co_argcount,
-                original_code.co_posonlyargcount,
-                original_code.co_kwonlyargcount,
-                original_code.co_nlocals,
-                original_code.co_stacksize,
-                original_code.co_flags,
-                original_code.co_code.replace(
-                    # Replace the hardcoded stride calculation with our dynamic one
-                    b"self.model.encoder.conv1.stride[0] * self.model.encoder.conv2.stride[0]",
-                    b"self.get_conv_stride()",
-                ),
-                original_code.co_consts,
-                original_code.co_names,
-                original_code.co_varnames,
-                original_code.co_filename,
-                original_code.co_name,
-                original_code.co_firstlineno,
-                original_code.co_lnotab,
-                original_code.co_freevars,
-                original_code.co_cellvars,
-            )
-            
-            # Create a new function with the modified code
-            new_generate = types.FunctionType(
-                modified_code,
-                original_generate.__globals__,
-                original_generate.__name__,
-                original_generate.__defaults__,
-                original_generate.__closure__,
-            )
-            
-            # Bind the function to the instance and call it
-            return new_generate(self, *args, **kwargs)
-        
-        # Add the stride calculation method to the mixin
-        WhisperGenerationMixin.get_conv_stride = get_conv_stride
-        # Replace the original generate method
-        WhisperGenerationMixin.generate = generate_wrapper
 
     
     def forward(
@@ -1862,9 +1803,70 @@ class CustomWhisperForConditionalGeneration(WhisperGenerationMixin, WhisperPreTr
         self.proj_out = nn.Linear(config.d_model, config.vocab_size, bias=False)
         self.max_target_positions = config.max_target_positions
 
+        self.patch_generate()
+
         # Initialize weights and apply final processing
         self.post_init()
 
+    # CUSTOM (Monkeypatch the generation method)
+    def patch_generate(self):
+        """
+        Monkey patches the WhisperGenerationMixin to use dynamic stride calculation
+        """
+        original_generate = WhisperGenerationMixin.generate
+        
+        def get_conv_stride(self):
+            """Calculate total stride of all conv layers"""
+            total_stride = 1
+            for layer in self.model.encoder.conv_layers:
+                total_stride *= layer.stride[0]
+            return total_stride
+        
+        def generate_wrapper(self, *args, **kwargs):
+            # Store the original function logic
+            original_code = original_generate.__code__
+            
+            # Create a modified version of the function that uses our stride calculation
+            modified_code = types.CodeType(
+                original_code.co_argcount,
+                original_code.co_posonlyargcount,
+                original_code.co_kwonlyargcount,
+                original_code.co_nlocals,
+                original_code.co_stacksize,
+                original_code.co_flags,
+                original_code.co_code.replace(
+                    # Replace the hardcoded stride calculation with our dynamic one
+                    b"self.model.encoder.conv1.stride[0] * self.model.encoder.conv2.stride[0]",
+                    b"self.get_conv_stride()",
+                ),
+                original_code.co_consts,
+                original_code.co_names,
+                original_code.co_varnames,
+                original_code.co_filename,
+                original_code.co_name,
+                original_code.co_firstlineno,
+                original_code.co_lnotab,
+                original_code.co_freevars,
+                original_code.co_cellvars,
+            )
+            
+            # Create a new function with the modified code
+            new_generate = types.FunctionType(
+                modified_code,
+                original_generate.__globals__,
+                original_generate.__name__,
+                original_generate.__defaults__,
+                original_generate.__closure__,
+            )
+            
+            # Bind the function to the instance and call it
+            return new_generate(self, *args, **kwargs)
+        
+        # Add the stride calculation method to the mixin
+        WhisperGenerationMixin.get_conv_stride = get_conv_stride
+        # Replace the original generate method
+        WhisperGenerationMixin.generate = generate_wrapper
+
     def get_encoder(self):
         return self.model.get_encoder()