This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,248 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
import importlib.util
|
||||
import os
|
||||
import requests
|
||||
from modelscope.hub.api import HubApi, ModelScopeConfig
|
||||
from modelscope.hub.utils.utils import get_cache_dir
|
||||
from pathlib import Path
|
||||
from tqdm import tqdm
|
||||
from typing import List, Optional
|
||||
|
||||
from .env import use_hf_hub
|
||||
from .logger import get_logger
|
||||
from .torch_utils import is_local_master, safe_ddp_context
|
||||
from .utils import subprocess_run
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def safe_snapshot_download(model_id_or_path: str,
|
||||
revision: Optional[str] = None,
|
||||
download_model: bool = True,
|
||||
use_hf: Optional[bool] = None,
|
||||
hub_token: Optional[str] = None,
|
||||
ignore_patterns: Optional[List[str]] = None,
|
||||
check_local: bool = False,
|
||||
**kwargs) -> str:
|
||||
"""Download model snapshot safely with DDP context protection.
|
||||
|
||||
This function attempts to download a model from HuggingFace or ModelScope hub,
|
||||
with support for local paths, subfolder specification, and distributed training
|
||||
context protection. It handles various path formats and provides flexible
|
||||
file filtering options.
|
||||
|
||||
Args:
|
||||
model_id_or_path (str): The model identifier on the hub (e.g., 'Qwen/Qwen2.5-7B-Instruct')
|
||||
or a local path to the model directory. Supports subfolder specification
|
||||
using colon syntax (e.g., 'model_id:subfolder').
|
||||
revision (Optional[str], optional): Specific model version/revision to download
|
||||
(branch name, tag, or commit hash). Defaults to None (latest version).
|
||||
download_model (bool, optional): Whether to download model weight files
|
||||
(.bin, .safetensors). If False, only config and tokenizer files are
|
||||
downloaded. Defaults to True.
|
||||
use_hf (Optional[bool], optional): Force using HuggingFace Hub (True) or ModelScope (False).
|
||||
If None, it is controlled by the environment variable `USE_HF`, which defaults to '0'.
|
||||
Default: None.
|
||||
hub_token (Optional[str], optional): Authentication token for accessing private
|
||||
or gated models. Defaults to None.
|
||||
ignore_patterns (Optional[List[str]], optional): List of glob patterns for files
|
||||
to exclude from download. If None, uses default patterns to exclude zip,
|
||||
gguf, pth, pt, and other auxiliary files. Defaults to None.
|
||||
check_local (bool, optional): Whether to check for a local directory matching
|
||||
the last component of model_id_or_path before attempting download.
|
||||
Defaults to False.
|
||||
**kwargs: Additional keyword arguments passed to the underlying hub download function.
|
||||
|
||||
Returns:
|
||||
str: Absolute path to the model directory where files are stored.
|
||||
|
||||
Raises:
|
||||
ValueError: If model_id_or_path starts with '/' (absolute path) and the path
|
||||
does not exist.
|
||||
Examples:
|
||||
>>> # Download from hub
|
||||
>>> model_dir = safe_snapshot_download('Qwen/Qwen2.5-7B-Instruct')
|
||||
|
||||
>>> # Download config only (no weights)
|
||||
>>> model_dir = safe_snapshot_download('Qwen/Qwen2.5-7B-Instruct', download_model=False)
|
||||
"""
|
||||
from swift.hub import get_hub
|
||||
if check_local:
|
||||
model_suffix = model_id_or_path.rsplit('/', 1)[-1]
|
||||
if os.path.exists(model_suffix):
|
||||
model_dir = os.path.abspath(os.path.expanduser(model_suffix))
|
||||
logger.info(f'Loading the model using local model_dir: {model_dir}')
|
||||
return model_dir
|
||||
if ignore_patterns is None:
|
||||
ignore_patterns = [
|
||||
'*.zip', '*.gguf', '*.pth', '*.pt', 'consolidated*', 'onnx/*', '*.safetensors.md', '*.msgpack', '*.onnx',
|
||||
'*.ot', '*.h5'
|
||||
]
|
||||
if not download_model:
|
||||
ignore_patterns += ['*.bin', '*.safetensors']
|
||||
hub = get_hub(use_hf)
|
||||
if model_id_or_path.startswith('~'):
|
||||
model_id_or_path = os.path.abspath(os.path.expanduser(model_id_or_path))
|
||||
model_path_to_check = '/'.join(model_id_or_path.split(':', 1))
|
||||
if os.path.exists(model_id_or_path):
|
||||
model_dir = model_id_or_path
|
||||
sub_folder = None
|
||||
elif os.path.exists(model_path_to_check):
|
||||
model_dir = model_path_to_check
|
||||
sub_folder = None
|
||||
else:
|
||||
if model_id_or_path.startswith('/'): # startswith
|
||||
raise ValueError(f"path: '{model_id_or_path}' not found")
|
||||
model_id_or_path = model_id_or_path.split(':', 1) # get sub_folder
|
||||
if len(model_id_or_path) == 1:
|
||||
model_id_or_path = [model_id_or_path[0], None]
|
||||
model_id_or_path, sub_folder = model_id_or_path
|
||||
if sub_folder is not None:
|
||||
kwargs['allow_patterns'] = [f"{sub_folder.rstrip('/')}/*"]
|
||||
with safe_ddp_context(hash_id=model_id_or_path):
|
||||
model_dir = hub.download_model(model_id_or_path, revision, ignore_patterns, token=hub_token, **kwargs)
|
||||
|
||||
logger.info(f'Loading the model using model_dir: {model_dir}')
|
||||
|
||||
model_dir = os.path.abspath(os.path.expanduser(model_dir))
|
||||
if sub_folder:
|
||||
model_dir = os.path.join(model_dir, sub_folder)
|
||||
assert os.path.isdir(model_dir), f'model_dir: {model_dir}'
|
||||
return model_dir
|
||||
|
||||
|
||||
def git_clone_github(github_url: str,
|
||||
*,
|
||||
local_repo_name: Optional[str] = None,
|
||||
branch: Optional[str] = None,
|
||||
commit_hash: Optional[str] = None) -> str:
|
||||
if github_url.endswith('.git'):
|
||||
github_url = github_url[:-4]
|
||||
git_cache_dir = os.path.join(get_cache_dir(), '_github')
|
||||
os.makedirs(git_cache_dir, exist_ok=True)
|
||||
if local_repo_name is None:
|
||||
github_url = github_url.rstrip('/')
|
||||
local_repo_name = github_url.rsplit('/', 1)[1]
|
||||
github_url = f'{github_url}.git'
|
||||
local_repo_path = os.path.join(git_cache_dir, local_repo_name)
|
||||
with safe_ddp_context('git_clone', use_barrier=True):
|
||||
repo_existed = os.path.exists(local_repo_path)
|
||||
if not is_local_master() and repo_existed:
|
||||
return local_repo_path
|
||||
if repo_existed:
|
||||
command = ['git', '-C', local_repo_path, 'fetch']
|
||||
subprocess_run(command)
|
||||
if branch is not None:
|
||||
command = ['git', '-C', local_repo_path, 'checkout', branch]
|
||||
subprocess_run(command)
|
||||
else:
|
||||
command = ['git', '-C', git_cache_dir, 'clone', github_url, local_repo_name]
|
||||
if branch is not None:
|
||||
command += ['--branch', branch]
|
||||
subprocess_run(command)
|
||||
|
||||
if commit_hash is not None:
|
||||
command = ['git', '-C', local_repo_path, 'reset', '--hard', commit_hash]
|
||||
subprocess_run(command)
|
||||
elif repo_existed:
|
||||
command = ['git', '-C', local_repo_path, 'pull']
|
||||
subprocess_run(command)
|
||||
logger.info(f'local_repo_path: {local_repo_path}')
|
||||
return local_repo_path
|
||||
|
||||
|
||||
def download_ms_file(url: str, local_path: str, cookies=None) -> None:
|
||||
if cookies is None:
|
||||
cookies = ModelScopeConfig.get_cookies()
|
||||
resp = requests.get(url, cookies=cookies, stream=True)
|
||||
with open(local_path, 'wb') as f:
|
||||
for data in tqdm(resp.iter_lines()):
|
||||
f.write(data)
|
||||
|
||||
|
||||
def _resolve_kernel_variant_str(repo_id: str) -> Optional[str]:
|
||||
"""Resolve the kernel build variant matching the current torch/cuda/platform
|
||||
by listing the ``build/`` folder of the ModelScope kernel repository. Returns
|
||||
``None`` if listing or parsing fails (caller should fall back to downloading
|
||||
the whole repo).
|
||||
"""
|
||||
try:
|
||||
from kernels.variants import parse_variant, resolve_variant
|
||||
files = HubApi().get_model_files(repo_id, root='build', recursive=False)
|
||||
variants = []
|
||||
for f in files:
|
||||
name = f.get('Name') or f.get('Path', '').rsplit('/', 1)[-1]
|
||||
if not name:
|
||||
continue
|
||||
try:
|
||||
variants.append(parse_variant(name))
|
||||
except ValueError:
|
||||
continue
|
||||
variant = resolve_variant(variants)
|
||||
return variant.variant_str if variant else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def patch_kernels() -> bool:
|
||||
"""Install a process-wide monkey patch on
|
||||
``transformers.integrations.hub_kernels.get_kernel`` so that kernel
|
||||
repositories are downloaded from ModelScope and loaded via
|
||||
``kernels.get_local_kernel``.
|
||||
|
||||
The runtime behavior is controlled by the ``USE_HF`` env (read on each
|
||||
``get_kernel`` call):
|
||||
- ``USE_HF=1``: fall back to the original HuggingFace-based loading.
|
||||
- otherwise (default): use ModelScope.
|
||||
|
||||
Returns True if the patch was installed, False if skipped (``kernels`` not
|
||||
installed, or the ``transformers`` integration is unavailable). Callers are
|
||||
expected to guarantee idempotency (e.g. via a module-level flag).
|
||||
"""
|
||||
if importlib.util.find_spec('kernels') is None:
|
||||
return False
|
||||
try:
|
||||
from kernels import get_local_kernel
|
||||
from transformers.integrations import hub_kernels
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
origin_get_kernel = hub_kernels.get_kernel
|
||||
|
||||
def patched_get_kernel(repo_id, *args, **kwargs):
|
||||
if use_hf_hub():
|
||||
return origin_get_kernel(repo_id, *args, **kwargs)
|
||||
try:
|
||||
variant_str = _resolve_kernel_variant_str(repo_id)
|
||||
allow_patterns = [f'build/{variant_str}/*'] if variant_str else None
|
||||
model_dir = safe_snapshot_download(repo_id, use_hf=False, allow_patterns=allow_patterns)
|
||||
package_name = repo_id.split('/')[-1].replace('-', '_')
|
||||
# kernels < 0.14
|
||||
kernel = get_local_kernel(Path(model_dir), package_name)
|
||||
logger.info(f'Loaded kernel `{repo_id}` from ModelScope: {model_dir}')
|
||||
return kernel
|
||||
except Exception as e:
|
||||
logger.warning(f'Failed to load kernel `{repo_id}` from ModelScope ({e}), fallback to HuggingFace.')
|
||||
return origin_get_kernel(repo_id, *args, **kwargs)
|
||||
|
||||
hub_kernels.get_kernel = patched_get_kernel
|
||||
return True
|
||||
|
||||
|
||||
def download_file(url: str) -> str:
|
||||
url = url.rstrip('/')
|
||||
file_name = url.rsplit('/', 1)[-1]
|
||||
cache_dir = os.path.join(get_cache_dir(), 'files')
|
||||
os.makedirs(cache_dir, exist_ok=True)
|
||||
file_path = os.path.join(cache_dir, file_name)
|
||||
if os.path.exists(file_path):
|
||||
return file_path
|
||||
resp = requests.get(url, stream=True)
|
||||
resp.raise_for_status()
|
||||
total_size = int(resp.headers.get('content-length', 0))
|
||||
with open(file_path, 'wb') as f, tqdm(
|
||||
total=total_size, unit='B', unit_scale=True, unit_divisor=1024, desc=file_name) as pbar:
|
||||
for chunk in resp.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
pbar.update(len(chunk))
|
||||
return file_path
|
||||
Reference in New Issue
Block a user