Upload HyenaDNAForCausalLM

modeling_hyena.py

@@ -349,8 +349,15 @@ class HyenaDNAPreTrainedModel(PreTrainedModel):
     supports_gradient_checkpointing = True
     _no_split_modules = ["HyenaBlock"]
     _skip_keys_device_placement = "past_key_values"
-
-    def _init_weights(self, initializer_range=0.02):
+    _keys_to_ignore_on_load_missing = [r"freq"]  # Shared tensors that safetensors merges
+
+    def _init_weights(self, module, initializer_range=0.02):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, std=initializer_range)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, std=initializer_range)
         # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
         #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
         #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
@@ -368,8 +375,8 @@ class HyenaDNAPreTrainedModel(PreTrainedModel):
 
 
 class HyenaDNAModel(HyenaDNAPreTrainedModel):
-    def __init__(self, config) -> None:
-        super().__init__(config)
+    def __init__(self, config, **kwargs) -> None:
+        super().__init__(config, **kwargs)
 
         self.backbone = HyenaLMBackbone(config)
         self.config = config
@@ -395,8 +402,8 @@ class HyenaDNAModel(HyenaDNAPreTrainedModel):
 
 class HyenaDNAForCausalLM(HyenaDNAPreTrainedModel):
 
-    def __init__(self, config):
-        super().__init__(config)
+    def __init__(self, config, **kwargs):
+        super().__init__(config, **kwargs)
         self.hyena = HyenaDNAModel(config)
         vocab_size = config.vocab_size
         if vocab_size % config.pad_vocab_size_multiple != 0:
@@ -476,9 +483,9 @@ class HyenaDNAForCausalLM(HyenaDNAPreTrainedModel):
 
 
 class HyenaDNAForSequenceClassification(HyenaDNAPreTrainedModel):
-    def __init__(self, config):
-        super().__init__(config)
-        self.num_labels = config.num_labels
+    def __init__(self, config, **kwargs):
+        super().__init__(config, **kwargs)
+        self.num_labels = kwargs.get("num_labels", config.num_labels)
         self.hyena = HyenaDNAModel(config)
         self.score = nn.Linear(config.d_model, self.num_labels, bias=False)