Better HF integration for MambaLMHeadModel

[Wauplin]

Hi @tridao and team,

This is a follow-up PR after #469 and in particular #469 (comment). This PR:

• adds huggingface_hub as an explicit dependency. It is already a dependency since mamba_ssm depends on transformers which depends on huggingface_hub but it's better to be explicit. I pinned a recent version that contains all recent PyTochModelHubMixin updates and fixes.
• removes PyTorchModelHubMixin from Mamba2 layer (introduced in Add HF integration, better discoverability #469)
• adds PyTorchModelHubMixin to MambaLMHeadModel for a better HF integration. I removed the existing from_pretrained / save_pretrained that were previously implemented. They still exists thanks to the mixin, and in a more robust way. The mixin also adds a push_to_hub method to directly save a model and push it to the Hub. All three helpers supports safetensors (if installed on the users's machine) and parameters like cache_dir/token/revision/etc. that can prove useful to users.

from mamba_ssm import MambaLMHeadModel

model = MambaLMHeadModel.from_pretrained("state-spaces/mamba2-130m")

(...)

model.push_to_hub("my-finetuned-mamba")

When doing MambaLMHeadModel(...).push_to_hub("username/my-cool-mamba"), a model card will be automatically created with some metadata in it (see docs). This improves a lot the UX on the Hub: better discoverability, better documentation, etc. In particular, I have added:

• library_name: mamba-ssm. I've opened a PR on our side to make mamba-ssm recognized as a library by the HF Hub (see Add MambaSSM as a library huggingface/huggingface.js#802). Users landing on a mamba_ssm model will automatically get a code snippet on how to instantiate the model + link to the mamba repo + download count enabled
• https://github.com/state-spaces/mamba as "repo_url" => all model cards will have a sentence "this model have been pushed using https://github.com/state-spaces/mamba" => better for documentation
• arXiv:2312.00752 and arXiv:2405.21060 as tags so that your papers will be automatically linked to all Mamba models uploaded on the Hub => better for referencing
• pipeline_tag: text-generation => models will appear when users filter models by task on the hub (https://huggingface.co/models?pipeline_tag=text-generation&sort=trending)

---

In parallel of this PR, it would be good to update existing models to add metadata as well. I opened 3 PRs to showcase what should be updated:

• https://huggingface.co/state-spaces/mamba-370m/discussions/3
• https://huggingface.co/state-spaces/mamba-130m/discussions/3
• https://huggingface.co/state-spaces/mamba2-130m/discussions/2

I have a script to open such a PR on other models. Let me know what you think and if you validate, I'll proceed with the others.

[Wauplin]

@tridao @NielsRogge I've update this PR + the description following your instructions. It is now ready to be reviewed. Only MambaLMHeadModel inherits from PyTorchModelHubMixin in the end, making it much simpler.

[NielsRogge]

Thanks, looks like a nice clean-up. Have you tested pushing and pulling a Mamba model to/from the hub (to ensure it works fine)?

[Wauplin]

Yes I did:

• Colab: https://colab.research.google.com/drive/1jDFxLmq0uWcQQrY2vGroTHbAIkkCfVYz?usp=sharing
• Pushed model: https://huggingface.co/Wauplin/my-cool-mamba