Allow saving / loading from Huggingface Hub preset #1510
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -27,8 +27,15 @@ | |
| except ImportError: | ||
| kagglehub = None | ||
|
|
||
| try: | ||
| import huggingface_hub | ||
| except ImportError: | ||
| huggingface_hub = None | ||
|
|
||
| KAGGLE_PREFIX = "kaggle://" | ||
| GS_PREFIX = "gs://" | ||
| HF_PREFIX = "hf://" | ||
|
|
||
| TOKENIZER_ASSET_DIR = "assets/tokenizer" | ||
| CONFIG_FILE = "config.json" | ||
| TOKENIZER_CONFIG_FILE = "tokenizer.json" | ||
|
|
@@ -69,6 +76,14 @@ def get_file(preset, path): | |
| url, | ||
| cache_subdir=os.path.join("models", subdir), | ||
| ) | ||
| elif preset.startswith(HF_PREFIX): | ||
| if huggingface_hub is None: | ||
| raise ImportError( | ||
| f"`from_preset()` requires the `huggingface_hub` package to load from '{preset}'. " | ||
| "Please install with `pip install huggingface_hub`." | ||
| ) | ||
| hf_handle = preset.removeprefix(HF_PREFIX) | ||
| return huggingface_hub.hf_hub_download(repo_id=hf_handle, filename=path) | ||
| elif os.path.exists(preset): | ||
| # Assume a local filepath. | ||
| return os.path.join(preset, path) | ||
mattdangerw marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
@@ -262,6 +277,15 @@ def upload_preset( | |
| if uri.startswith(KAGGLE_PREFIX): | ||
| kaggle_handle = uri.removeprefix(KAGGLE_PREFIX) | ||
| kagglehub.model_upload(kaggle_handle, preset) | ||
| elif uri.startswith(HF_PREFIX): | ||
| if huggingface_hub is None: | ||
| raise ImportError( | ||
| f"`upload_preset()` requires the `huggingface_hub` package to upload to '{uri}'. " | ||
| "Please install with `pip install huggingface_hub`." | ||
| ) | ||
| hf_handle = uri.removeprefix(HF_PREFIX) | ||
| repo_url = huggingface_hub.create_repo(repo_id=hf_handle, exist_ok=True) | ||
| huggingface_hub.upload_folder(repo_id=repo_url.repo_id, folder_path=preset) | ||
| else: | ||
| raise ValueError( | ||
| f"Unexpected URI `'{uri}'`. Kaggle upload format should follow " | ||
|
There was a problem hiding this comment. Same here, might want to reword this error message. Side note: we are kinda inconsistent in how we refer to these model handles. I'm not sure if we should call these URIs or something else, but we should be consistent in our wording. No need to fix on this PR. There was a problem hiding this comment. I've update the error message in 24ef262. I kept the |
||
|
|
||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
What do the error messages look like if a handle is unformed? Will it read well enough, or should we validate roughly here so we can have a message similar to the Kaggle error above?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
At this point,
hf_handlemust correspond to a repo_id so something in the formusername/repo_name(e.g. "google/gemma-7b") since the hf prefix has been removed. If it's not the case, anHFValidationError(which is a customValueError) is raised. Here are the validation rules we are checking.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I just pushed 24ef262 to raise a ValueError similar to the kaggle handle one. I tried to be as consistent as possible. Please let me know what you think :)