Update configuration_pebble.py

configuration_pebble.py

@@ -6,35 +6,36 @@ class PebbleConfig(PretrainedConfig):
     def __init__(
         self,
         vocab_size=2048,
-        d_model=608,
-        n_blocks=8,
-        n_heads=8,
-        att_hidden=2432,
-        rope_theta=10000.0,
-        seq_len=2048,
-        mamba_d_state=128,
-        mamba_d_conv=4,
-        mamba_expand=2,
-        mamba_headdim=64,
-        block_pattern="3:1 mamba:attn",
-        bos_token_id=1,
-        eos_token_id=1,
+        hidden_size=608,
+        intermediate_size=2432,
+        num_hidden_layers=8,
+        num_attention_heads=8,
+        block_pattern="mmma|mmma",
+        hybrid_ratio="3:1 mamba2:attention",
+        max_position_embeddings=2048,
+        rms_norm_eps=1e-6,
+        tie_word_embeddings=True,
+        mamba2=None,
+        attention=None,
         **kwargs,
     ):
         self.vocab_size = vocab_size
-        self.d_model = d_model
-        self.n_blocks = n_blocks
-        self.n_heads = n_heads
-        self.att_hidden = att_hidden
-        self.rope_theta = rope_theta
-        self.seq_len = seq_len
-        self.mamba_d_state = mamba_d_state
-        self.mamba_d_conv = mamba_d_conv
-        self.mamba_expand = mamba_expand
-        self.mamba_headdim = mamba_headdim
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
         self.block_pattern = block_pattern
-        super().__init__(
-            bos_token_id=bos_token_id,
-            eos_token_id=eos_token_id,
-            **kwargs,
-        )
+        self.hybrid_ratio = hybrid_ratio
+        self.max_position_embeddings = max_position_embeddings
+        self.rms_norm_eps = rms_norm_eps
+        self.tie_word_embeddings = tie_word_embeddings
+        
+        # Default dictionaries if not provided in config.json
+        self.mamba2 = mamba2 or {
+            "d_state": 128, "d_conv": 4, "expand": 2, 
+            "headdim": 64, "use_mem_eff_path": True
+        }
+        self.attention = attention or {
+            "rope_theta": 10000.0, "is_causal": True
+        }
+        super().__init__(**kwargs)