Implement a FlaxModelHubMixin

[Wauplin]

Related to this slack thread (private) cc @NielsRogge @osanseviero

We have a generic ModelHubMixin that helps integrating any machine learning framework with the Hub as long as one inherits an object from both the framework's main class and ModelHubMixin. Once inherited, one have to implement the _from_pretrained and _save_pretrained methods to specify how the files should be serialized on the Hub. From that mixin, we derived a PyTorchModelHubMixin that is specifically targeted for PyTorch. The serialization/deserialization is already implemented, making it directly plug-able to any Pytorch model. Here is a more thorough description in the documentation.

A good addition would be to implement a FlaxModelHubMixin class directly in huggingface_hub, in the same way PyTorchModelHubMixin is already implemented. The purpose is to be able to host/load any Flax models on the Hugging Face Hub.

EDIT: tagging it as good first issue Good for newcomers as it is a self-contained issue that can be implemented on its own. It requires a bit of knowledge about Flax serialization methods. Feel free to ask questions.

[sanjeevk-os]

Hi @Wauplin : Can I work on this issue ?

[Wauplin]

Hi @sanjeevk-os thanks for proposing to work on this PR! ❤️ Yes, it is free to work on so I can assign it to you. Please let me know if you need any assistance.
I would suggest to start opening a PR with only the class implementation itself. Once we've tested and agreed on the logic itself, you could continue for a proper PR with implementation + docs + unit tests. Saying this since this issue is about defining a new class so the interface might require some discussion. In any case, feel free to start and thanks again for the proposition!

[lappemic]

Hey @Wauplin, i would love to contribute here. I have to say i do not have any experience in Flax though and would need some guidance. 😄

[Wauplin]

Hi @lappemic, thanks for proposing your help on this one. However I do think it's best if the mixin gets implemented by someone with at least some knowledge of the Flax ecosystem, at least to know the different corner cases that could exist. I wouldn't want to do it myself either to be honest (I have set the "good first issue" label because it is an issue that requires little knowledge of huggingface_hub).

Would you be interested in other pure-Python pure-huggingface_hub good first issues like #1595 or #1808 maybe?

[lappemic]

Thanks a lot for your open reply and pointing to more suitable issues @Wauplin. I would love to work on #1808. I am having a look at it.