keras-team / keras-nlp

Modular Natural Language Processing workflows with Keras
Apache License 2.0
758 stars 227 forks source link

Allow saving / loading from Huggingface Hub preset #1510

Closed Wauplin closed 5 months ago

Wauplin commented 5 months ago

Solves https://github.com/keras-team/keras-nlp/issues/1294. As mentioned in https://github.com/keras-team/keras-nlp/issues/1294#issuecomment-1966864503, this PR adds support for the hf:// prefix to load presets from the Huggingface Hub.

The integration requires the huggingface_hub library. Authentication can be configured with the HF_TOKEN environment variable (only for private models or for uploads, similarly to KaggleHub). Here is a Colab notebook showcasing it.

import keras_nlp
from keras_nlp.models import BertClassifier
from keras_nlp.utils.preset_utils import save_to_preset, upload_preset

classifier = BertClassifier.from_preset("bert_base_en_uncased")
(...) # train/retrain/fine-tune

# Save to local folder
save_to_preset(classifier, "bert_base_en_uncased_retrained")

# Upload to Hugging Face Hub
upload_preset("hf://Wauplin/bert_base_en_uncased_retrained", "bert_base_en_uncased_retrained", allow_incomplete=True)

# Reload from the Hugging Face Hub
classifier_reloaded = keras_nlp.models.BertClassifier.from_preset(
    "hf://Wauplin/bert_base_en_uncased_retrained",
    num_classes=2,
    activation="softmax",
)

Here is how it looks like once uploaded on the Hub: https://huggingface.co/Wauplin/bert_base_en_uncased_retrained/tree/main.. If we go this way, I think we should also upload a default model card with keras-nlp tag to make all KerasNLP models discoverable on the Hub. On the Hugging Face side, we could make KerasNLP an official library (e.g. searchable, with code snippets, download counts, etc.).

In the current implementation, saving to "hf://Wauplin/bert_base_en_uncased_retrained" will save the model locally to Wauplin/bert_base_en_uncased_retrained subfolder + create the repository on the Hub + upload the local folder to this repo on the Hub. An alternative could be to save to a temporary folder before uploading to the Hub (to avoid the local copy). Both solutions are correct in my opinion, it's more a matter of how the KerasNLP envision the save_to_preset method.

mattdangerw commented 5 months ago

@Wauplin with https://github.com/keras-team/keras-nlp/pull/1512 I think we are all ready to go here.

Note that we will keep landing features related to the whole upload flow (in particular uploading a classifier with extra head weights, uploading lora weights that are essentially a diff on another preset). But we will keep the upload_preset flow, so I think this is unblocked.

Wauplin commented 5 months ago

Thanks for the ping @mattdangerw! I've just updated the PR accordingly and we should now be good to go :) And agree with you, we shouldn't have much HF-specific logic apart from that given how isolated upload_preset and get_file are.

Regarding documentation, do I have to update some markdown somewhere or are the different preset handlers not really documented for now? Please let me know if I can be of any assistance here.

SamanehSaadat commented 5 months ago

Thanks for the PR, @Wauplin! Looks great! Regarding documentation, for now, you can update the upload_preset() docstring here to include info about hf://. I'm working on preparing a Kaggle upload guide and I'll make sure to include info about HuggingFace as well.

Wauplin commented 5 months ago

Thanks both for the review and feedback! I have addressed all comments and completed the upload_preset docstring to mention hf://. I did not add it to get_file since kaggle preset was not documented there either but happy to document them both if you think it makes sense.

Wauplin commented 5 months ago

Great thanks for the approval! I just pushed a commit to fix linting. Hope it's fine now :)

SamanehSaadat commented 5 months ago

Hi @Wauplin!

Thanks again for the PR! I was trying to create a demo for HF upload and realized when you want to upload a model, you need to create the model on the HF web UI first. upload_folder() documentation mentions that the folder can be uploaded to an existing repo.

Just wanted to make sure my understanding is correct and there isn't any way to upload a model folder if the model hasn't been created on the HF website.

Wauplin commented 5 months ago

I was trying to create a demo for HF upload and realized when you want to upload a model, you need to create the model on the HF web UI first. upload_folder() documentation mentions that the folder can be uploaded to an existing repo.

@SamanehSaadat yes that's true but you can use create_repo to create the repo on the Hub first. This is actually what we are doing here.