mirror of
https://github.com/modelscope/modelscope.git
synced 2026-07-12 21:39:26 +02:00
284 lines
11 KiB
Python
284 lines
11 KiB
Python
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import logging
|
|
from argparse import ArgumentParser
|
|
|
|
from modelscope.cli.base import CLICommand
|
|
from modelscope.cli.utils import concurrent_download
|
|
from modelscope.hub.api import HubApi
|
|
from modelscope.hub.constants import DEFAULT_MAX_WORKERS, DEFAULT_SKILLS_DIR
|
|
from modelscope.hub.file_download import (dataset_file_download,
|
|
model_file_download)
|
|
from modelscope.hub.snapshot_download import (dataset_snapshot_download,
|
|
snapshot_download)
|
|
from modelscope.hub.utils.utils import convert_patterns, resolve_endpoint
|
|
from modelscope.utils.constant import DEFAULT_DATASET_REVISION
|
|
from modelscope.utils.logger import get_logger
|
|
|
|
logger = get_logger(log_level=logging.WARNING)
|
|
|
|
|
|
def subparser_func(args):
|
|
""" Function which will be called for a specific sub parser.
|
|
"""
|
|
return DownloadCMD(args)
|
|
|
|
|
|
class DownloadCMD(CLICommand):
|
|
name = 'download'
|
|
|
|
def __init__(self, args):
|
|
self.args = args
|
|
|
|
@staticmethod
|
|
def define_args(parsers: ArgumentParser):
|
|
""" define args for download command.
|
|
"""
|
|
parser: ArgumentParser = parsers.add_parser(DownloadCMD.name)
|
|
group = parser.add_mutually_exclusive_group()
|
|
group.add_argument(
|
|
'--model',
|
|
type=str,
|
|
help='The id of the model to be downloaded. For download, '
|
|
'the id of either a model or dataset must be provided.')
|
|
group.add_argument(
|
|
'--dataset',
|
|
type=str,
|
|
help='The id of the dataset to be downloaded. For download, '
|
|
'the id of either a model or dataset must be provided.')
|
|
group.add_argument(
|
|
'--collection',
|
|
type=str,
|
|
default=None,
|
|
help='The ID of the collection to download (skills only)')
|
|
parser.add_argument(
|
|
'repo_id',
|
|
type=str,
|
|
nargs='?',
|
|
default=None,
|
|
help='Optional, '
|
|
'ID of the repo to download, It can also be set by --model or --dataset.'
|
|
)
|
|
parser.add_argument(
|
|
'--repo-type',
|
|
choices=['model', 'dataset'],
|
|
default='model',
|
|
help="Type of repo to download from (defaults to 'model').",
|
|
)
|
|
parser.add_argument(
|
|
'--token',
|
|
type=str,
|
|
default=None,
|
|
help='Optional. Access token to download controlled entities.')
|
|
parser.add_argument(
|
|
'--revision',
|
|
type=str,
|
|
default=None,
|
|
help='Revision of the entity (e.g., model).')
|
|
parser.add_argument(
|
|
'--cache_dir',
|
|
type=str,
|
|
default=None,
|
|
help='Cache directory to save entity (e.g., model).')
|
|
parser.add_argument(
|
|
'--local_dir',
|
|
type=str,
|
|
default=None,
|
|
help='File will be downloaded to local location specified by'
|
|
'local_dir, in this case, cache_dir parameter will be ignored.')
|
|
parser.add_argument(
|
|
'files',
|
|
type=str,
|
|
default=None,
|
|
nargs='*',
|
|
help='Specify relative path to the repository file(s) to download.'
|
|
"(e.g 'tokenizer.json', 'onnx/decoder_model.onnx').")
|
|
parser.add_argument(
|
|
'--include',
|
|
nargs='*',
|
|
default=None,
|
|
type=str,
|
|
help='Glob patterns to match files to download.'
|
|
'Ignored if file is specified')
|
|
parser.add_argument(
|
|
'--exclude',
|
|
nargs='*',
|
|
type=str,
|
|
default=None,
|
|
help='Glob patterns to exclude from files to download.'
|
|
'Ignored if file is specified')
|
|
parser.add_argument(
|
|
'--endpoint',
|
|
type=str,
|
|
default=None,
|
|
help='ModelScope server endpoint, e.g. modelscope.cn or '
|
|
'modelscope.ai Full URL like '
|
|
'https://modelscope.cn is also accepted. Scheme (https://) is '
|
|
'auto-completed if omitted. Falls back to env MODELSCOPE_DOMAIN, '
|
|
'then defaults to https://www.modelscope.cn. '
|
|
'When omitted, the CLI auto-detects the correct site '
|
|
'(cn/intl) for download.')
|
|
parser.add_argument(
|
|
'--max-workers',
|
|
type=int,
|
|
default=DEFAULT_MAX_WORKERS,
|
|
help='The maximum number of workers to download files.')
|
|
|
|
parser.set_defaults(func=subparser_func)
|
|
|
|
def execute(self):
|
|
if self.args.model or self.args.dataset:
|
|
# the position argument of files will be put to repo_id.
|
|
if self.args.repo_id is not None:
|
|
if self.args.files:
|
|
self.args.files.insert(0, self.args.repo_id)
|
|
else:
|
|
self.args.files = [self.args.repo_id]
|
|
else:
|
|
if self.args.repo_id is not None:
|
|
if self.args.repo_type == 'model':
|
|
self.args.model = self.args.repo_id
|
|
elif self.args.repo_type == 'dataset':
|
|
self.args.dataset = self.args.repo_id
|
|
else:
|
|
raise Exception('Not support repo-type: %s'
|
|
% self.args.repo_type)
|
|
if not self.args.model and not self.args.dataset and not self.args.collection:
|
|
raise Exception('Model, dataset, or collection must be set.')
|
|
if self.args.endpoint:
|
|
endpoint = resolve_endpoint(self.args.endpoint)
|
|
else:
|
|
endpoint = None
|
|
cookies = None
|
|
if self.args.token is not None:
|
|
api = HubApi(endpoint=endpoint)
|
|
cookies = api.get_cookies(access_token=self.args.token)
|
|
if self.args.model:
|
|
if len(self.args.files) == 1: # download single file
|
|
model_file_download(
|
|
self.args.model,
|
|
self.args.files[0],
|
|
cache_dir=self.args.cache_dir,
|
|
local_dir=self.args.local_dir,
|
|
revision=self.args.revision,
|
|
cookies=cookies,
|
|
token=self.args.token,
|
|
endpoint=endpoint)
|
|
elif len(
|
|
self.args.files) > 1: # download specified multiple files.
|
|
snapshot_download(
|
|
self.args.model,
|
|
revision=self.args.revision,
|
|
cache_dir=self.args.cache_dir,
|
|
local_dir=self.args.local_dir,
|
|
allow_file_pattern=self.args.files,
|
|
max_workers=self.args.max_workers,
|
|
cookies=cookies,
|
|
token=self.args.token,
|
|
endpoint=endpoint)
|
|
else: # download repo
|
|
snapshot_download(
|
|
self.args.model,
|
|
revision=self.args.revision,
|
|
cache_dir=self.args.cache_dir,
|
|
local_dir=self.args.local_dir,
|
|
allow_file_pattern=convert_patterns(self.args.include),
|
|
ignore_file_pattern=convert_patterns(self.args.exclude),
|
|
max_workers=self.args.max_workers,
|
|
cookies=cookies,
|
|
token=self.args.token,
|
|
endpoint=endpoint)
|
|
print(f'\nSuccessfully Downloaded from model {self.args.model}.\n')
|
|
elif self.args.dataset:
|
|
dataset_revision: str = self.args.revision if self.args.revision else DEFAULT_DATASET_REVISION
|
|
if len(self.args.files) == 1: # download single file
|
|
dataset_file_download(
|
|
self.args.dataset,
|
|
self.args.files[0],
|
|
cache_dir=self.args.cache_dir,
|
|
local_dir=self.args.local_dir,
|
|
revision=dataset_revision,
|
|
cookies=cookies,
|
|
token=self.args.token,
|
|
endpoint=endpoint)
|
|
elif len(
|
|
self.args.files) > 1: # download specified multiple files.
|
|
dataset_snapshot_download(
|
|
self.args.dataset,
|
|
revision=dataset_revision,
|
|
cache_dir=self.args.cache_dir,
|
|
local_dir=self.args.local_dir,
|
|
allow_file_pattern=self.args.files,
|
|
max_workers=self.args.max_workers,
|
|
cookies=cookies,
|
|
token=self.args.token,
|
|
endpoint=endpoint)
|
|
else: # download repo
|
|
dataset_snapshot_download(
|
|
self.args.dataset,
|
|
revision=dataset_revision,
|
|
cache_dir=self.args.cache_dir,
|
|
local_dir=self.args.local_dir,
|
|
allow_file_pattern=convert_patterns(self.args.include),
|
|
ignore_file_pattern=convert_patterns(self.args.exclude),
|
|
max_workers=self.args.max_workers,
|
|
cookies=cookies,
|
|
token=self.args.token,
|
|
endpoint=endpoint)
|
|
print(
|
|
f'\nSuccessfully Downloaded from dataset {self.args.dataset}.\n'
|
|
)
|
|
elif self.args.collection:
|
|
api = HubApi(endpoint=endpoint, token=self.args.token)
|
|
local_dir = self.args.local_dir or DEFAULT_SKILLS_DIR
|
|
data = api.get_collection(
|
|
self.args.collection, repo_type='skill', endpoint=endpoint)
|
|
elements = data.get('CollectionElements',
|
|
{}).get('CollectionElementVoList', [])
|
|
|
|
logger.info(
|
|
f'Collection {self.args.collection} has {len(elements)} elements.'
|
|
)
|
|
|
|
if not elements:
|
|
print(f'No skill elements found in collection: '
|
|
f'{self.args.collection}')
|
|
return
|
|
|
|
# Validate elements have required fields
|
|
valid_elements = []
|
|
for elem in elements:
|
|
if not elem.get('ElementPath') or not elem.get('ElementName'):
|
|
logger.warning('Skipping malformed collection element: %s',
|
|
elem)
|
|
continue
|
|
valid_elements.append(elem)
|
|
|
|
if not valid_elements:
|
|
print(f'No valid skill elements found in collection: '
|
|
f'{self.args.collection}')
|
|
return
|
|
|
|
print(f'Found {len(valid_elements)} skill(s) in collection, '
|
|
f'downloading...')
|
|
|
|
def _download_one_skill(element):
|
|
element_path = element['ElementPath']
|
|
element_name = element['ElementName']
|
|
skill_id = f'{element_path}/{element_name}'
|
|
try:
|
|
skill_dir = api.download_skill(
|
|
skill_id=skill_id,
|
|
local_dir=local_dir,
|
|
endpoint=endpoint)
|
|
return (skill_id, skill_dir, None)
|
|
except Exception as e:
|
|
return (skill_id, None, str(e))
|
|
|
|
concurrent_download(
|
|
_download_one_skill,
|
|
valid_elements,
|
|
max_workers=self.args.max_workers,
|
|
item_name='skill')
|
|
else:
|
|
pass # noop
|