zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-30 02:53:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

update snapshot_download

Browse Source
This commit is contained in:
xingjun.wang 2025-01-06 00:40:25 +08:00
parent 5c535a28c0
commit 2954fc56dd
1 changed files with 33 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
llms/mlx_lm/utils.py
View File

@ -16,11 +16,13 @@ from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Type,
import mlx.core as mx import mlx.core as mx
import mlx.nn as nn import mlx.nn as nn
if os.getenv('MLX_USE_MODELSCOPE', 'False').lower() == 'true':
print(">> Using ModelScope") use_modelscope = os.getenv('MLX_USE_MODELSCOPE', 'False').lower() == 'true'
if use_modelscope:
from modelscope import snapshot_download from modelscope import snapshot_download
else: else:
from huggingface_hub import snapshot_download from huggingface_hub import snapshot_download
from mlx.utils import tree_flatten, tree_reduce from mlx.utils import tree_flatten, tree_reduce
from transformers import PreTrainedTokenizer from transformers import PreTrainedTokenizer
@ -158,12 +160,25 @@ def get_model_path(path_or_hf_repo: str, revision: Optional[str] = None) -> Path
Path: The path to the model. Path: The path to the model.
""" """
model_path = Path(path_or_hf_repo) model_path = Path(path_or_hf_repo)
print(f'>>model_path: {model_path}')
revision = revision or 'master'
print(f'>>revision: {revision}')
if not model_path.exists(): if not model_path.exists():
try: try:
if use_modelscope:
model_path = Path(
snapshot_download(
model_id=path_or_hf_repo,
revision=revision or 'master',
allow_patterns=[
"*.json",
"*.safetensors",
"*.py",
"tokenizer.model",
"*.tiktoken",
"*.txt",
],
)
)
else:
model_path = Path( model_path = Path(
snapshot_download( snapshot_download(
repo_id=path_or_hf_repo, repo_id=path_or_hf_repo,