update script

modeling_granite.py

@@ -1,4 +1,6 @@
+import math
 import numbers
+import warnings
 from enum import Enum
 from typing import Optional, Tuple, Union
 
@@ -846,7 +848,7 @@ class GranitePreTrainedModel(PreTrainedModel):
         self.initializer_range = config.initializer_range
 
     def _init_weights(self, module: nn.Module) -> None:
-        if isinstance(module, (nn.LayerNorm, RMSNorm, RoPE)):
+        if isinstance(module, (nn.LayerNorm, RMSNorm, Alibi, RoPE)):
             module.reset_parameters()
         elif isinstance(module, nn.Linear):
             nn.init.normal_(module.weight, mean=0, std=self.initializer_range)
@@ -1104,15 +1106,15 @@ class GraniteModel(GranitePreTrainedModel):
 
     def _prepare_a_bunch_of_stuff(
         self,
-        input_ids: torch.Tensor = None,
-        past_key_values: DynamicCache = None,
-        attention_mask: torch.Tensor = None,
-        token_type_ids: torch.Tensor = None,
-        position_ids: torch.Tensor = None,
-        inputs_embeds: torch.Tensor = None,
-        use_cache: bool = None,
-        output_hidden_states: bool = None,
-        return_dict: bool = None,
+        input_ids: torch.Tensor,
+        past_key_values: DynamicCache,
+        attention_mask: torch.Tensor,
+        token_type_ids: torch.Tensor,
+        position_ids: torch.Tensor,
+        inputs_embeds: torch.Tensor,
+        use_cache: bool,
+        output_hidden_states: bool,
+        return_dict: bool,
     ) -> Tuple[
         bool,
         bool,