Upload chess-gpt-eval/mamba_lm.py with huggingface_hub

chess-gpt-eval/mamba_lm.py

@@ -5,7 +5,8 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from mamba import Mamba, MambaConfig, RMSNorm
+from mamba import MambaConfig #Mamba, MambaConfig, RMSNorm
+from mamba_ssm import MambaLMHeadModel
 
 """
 
@@ -65,7 +66,7 @@ def from_pretrained(name: str):
     config_data = load_config_hf(name)
     config = MambaLMConfig(d_model=config_data['d_model'], n_layers=config_data['n_layer'], vocab_size=config_data['vocab_size'])
 
-    model = MambaLM(config)
+    model = MambaLMHeadModel(config)
 
     # copy weights
     state_dict = load_state_dict_hf(name)
@@ -87,7 +88,7 @@ class MambaLM(nn.Module):
     def __init__(self, lm_config: MambaLMConfig):
         super().__init__()
         self.lm_config = lm_config
-        self.config = lm_config.to_mamba_config()
+        self.config = lm_config#.to_mamba_config()
 
         self.embedding = nn.Embedding(self.lm_config.vocab_size, self.config.d_model)
         self.mamba = Mamba(self.config)