Update configuration_mistral.py

configuration_mistral.py

@@ -131,6 +131,18 @@ class MistralConfig(PretrainedConfig):
         rope_scaling=None,
         rope_theta=10000.0,
         sliding_window=4096,
+        attention_dropout=0.0,
+        max_thoughts=16,
+        max_temperature=10,
+        merged_talk_heads=True,
+        merged_lm_and_talk_heads=False,
+        merged_lm_and_think_heads=True,
+        use_concat_talk_head=True,
+        use_shallow_think=True,
+        use_shallow_talk=False,
+        use_complex_think_head=False,
+        use_complex_talk_head=True,
+        use_weighted_talk_head=True,
         **kwargs,
     ):
         self.vocab_size = vocab_size
@@ -140,7 +152,19 @@ class MistralConfig(PretrainedConfig):
         self.num_hidden_layers = num_hidden_layers
         self.num_attention_heads = num_attention_heads
         self.sliding_window = sliding_window
-
+        attention_dropout=0.0,
+        max_thoughts=16,
+        max_temperature=10,
+        complexity_factor = 0.5,
+        merged_talk_heads=True,
+        merged_lm_and_talk_heads=False,
+        merged_lm_and_think_heads=True,
+        use_concat_talk_head=True,
+        use_shallow_think=True,
+        use_shallow_talk=False,
+        use_complex_think_head=False,
+        use_complex_talk_head=True,
+        use_weighted_talk_head=True,
         # for backward compatibility
         if num_key_value_heads is None:
             num_key_value_heads = num_attention_heads
@@ -153,7 +177,22 @@ class MistralConfig(PretrainedConfig):
         self.rope_scaling = rope_scaling
         self._rope_scaling_validation()
         self.rope_theta = rope_theta
+        self.attention_dropout = attention_dropout
+        self.max_thoughts = max_thoughts
+        self.complexity_factor = complexity_factor
+        self.max_temperature = max_temperature
+        self.merged_talk_heads = merged_talk_heads
+        self.merged_lm_and_talk_heads = merged_lm_and_talk_heads
+        self.merged_lm_and_think_heads = merged_lm_and_think_heads
+        self.use_concat_talk_head = use_concat_talk_head
+        self.use_shallow_think = use_shallow_think
+        self.use_shallow_talk = use_shallow_talk
+        self.use_complex_think_head = use_complex_think_head
+        self.use_complex_talk_head = use_complex_talk_head
+        self.use_weighted_talk_head = use_weighted_talk_head
+        self.hidden_dropout_prob = hidden_dropout_prob
 
+        
         super().__init__(
             pad_token_id=pad_token_id,
             bos_token_id=bos_token_id,