Unable to run prompt_mamba_with_context.py

[quaid281]

I keep getting this error when running prompt_mamba_with_context.py

Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
Traceback (most recent call last):
File "/home/ec2-user/SageMaker/mamba-dive/prompt_mamba_with_context.py", line 42, in
model = MambaLMHeadModel.from_pretrained(model, device="cuda", dtype=torch.float16)
File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/mamba_ssm/models/mixer_seq_simple.py", line 245, in from_pretrained
model.load_state_dict(load_state_dict_hf(pretrained_model_name, device=device, dtype=dtype))
File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/mamba_ssm/utils/hf.py", line 18, in load_state_dict_hf
return torch.load(resolved_archive_file, map_location=mapped_device)
File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/serialization.py", line 815, in load
return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/serialization.py", line 1033, in _legacy_load
magic_number = pickle_module.load(f, **pickle_load_args)
_pickle.UnpicklingError: invalid load key, 'v'.

[gschoeni]

Are you running it with the model from Oxen-AI/mamba-130m-context?

python prompt_mamba_with_context.py Oxen-AI/mamba-130m-context

Also what are your pytorch / transformers versions? Maybe you can dump a pip freeze and we can see what's different.

[quaid281]

requirements.txt

I'm running it with the model from Oxen-AI/mamba-130m-context in an AWS Sagemaker Jupyter notebook. environment

Here is my requirements.txt