add mamba dataclass args

Author: jmercat

open_lm/model.py

@@ -100,6 +100,7 @@ class Params:
     positional_embedding_type: str = "rotary"
     ffn_type: str = "swiglu"
 
+
 @dataclass
 class MambaParams:
     d_model: int = 2560
@@ -114,6 +115,7 @@ class MambaParams:
     tie_embeddings: bool = True
     weight_tying: bool = False
 
+
 def get_pos_embed(args: Params):
     head_dim = args.dim // args.n_heads
     if args.positional_embedding_type == "rotary":

open_lm/utils/transformers/hf_config.py

@@ -40,5 +40,5 @@ def __init__(
 
     def set_params(self, params: Params):
         self.tie_word_embeddings = params.weight_tying
-        for field in fields(Params):
+        for field in fields(params):
             setattr(self, field.name, getattr(params, field.name))