Update model.py

model.py

@@ -271,7 +271,7 @@ class WhisperAttention(nn.Module):
         bias: bool = True,
         is_causal: bool = False,
         layer_idx: Optional[int] = None,
-        config: Optional[WhisperConfig] = None,
+        config: Optional[CustomWhisperConfig] = None,
     ):
         super().__init__()
         self.embed_dim = embed_dim
@@ -616,7 +616,7 @@ WHISPER_ATTENTION_CLASSES = {
 
 # Copied from transformers.models.mbart.modeling_mbart.MBartEncoderLayer with MBart->Whisper, MBART->WHISPER
 class WhisperEncoderLayer(nn.Module):
-    def __init__(self, config: WhisperConfig):
+    def __init__(self, config: CustomWhisperConfig):
         super().__init__()
         self.embed_dim = config.d_model
 
@@ -686,7 +686,7 @@ class WhisperEncoderLayer(nn.Module):
 
 
 class WhisperDecoderLayer(nn.Module):
-    def __init__(self, config: WhisperConfig, layer_idx: int = None):
+    def __init__(self, config: CustomWhisperConfig, layer_idx: int = None):
         super().__init__()
         self.embed_dim = config.d_model
 
@@ -803,7 +803,7 @@ class WhisperDecoderLayer(nn.Module):
 
 
 class WhisperPreTrainedModel(PreTrainedModel):
-    config_class = WhisperConfig
+    config_class = CustomWhisperConfig
     base_model_prefix = "model"
     main_input_name = "input_features"
     supports_gradient_checkpointing = True
@@ -982,7 +982,7 @@ class WhisperEncoder(WhisperPreTrainedModel):
         config: WhisperConfig
     """
 
-    def __init__(self, config: WhisperConfig):
+    def __init__(self, config: CustomWhisperConfig):
         super().__init__(config)
         self.dropout = config.dropout
         self.layerdrop = config.encoder_layerdrop
@@ -1271,7 +1271,7 @@ class WhisperDecoder(WhisperPreTrainedModel):
 
     main_input_name = "input_ids"
 
-    def __init__(self, config: WhisperConfig):
+    def __init__(self, config: CustomWhisperConfig):
         super().__init__(config)
         self.dropout = config.dropout
         self.layerdrop = config.decoder_layerdrop
@@ -1674,7 +1674,7 @@ class WhisperDecoder(WhisperPreTrainedModel):
     WHISPER_START_DOCSTRING,
 )
 class WhisperModel(WhisperPreTrainedModel):
-    def __init__(self, config: WhisperConfig):
+    def __init__(self, config: CustomWhisperConfig):
         super().__init__(config)
 
         self.encoder = WhisperEncoder(config)
@@ -1849,7 +1849,7 @@ class CustomWhisperForConditionalGeneration(WhisperGenerationMixin, WhisperPreTr
     base_model_prefix = "model"
     _tied_weights_keys = ["proj_out.weight"]
 
-    def __init__(self, config: WhisperConfig):
+    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)