Skip to content

Commit 63ff6ec

Browse files
wangxingjun778wangxingjun778
authored
Addgated_mode for create repo (#1730)
1 parent 491e3df commit 63ff6ec

3 files changed

Lines changed: 74 additions & 8 deletions

File tree

‎modelscope/cli/create.py‎

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@
33

44
from modelscope.cli.base import CLICommand
55
from modelscope.hub.api import HubApi
6-
from modelscope.hub.constants import (Licenses, ModelVisibility, Visibility,
7-
VisibilityMap)
6+
from modelscope.hub.constants import (GatedMode, Licenses, ModelVisibility,
7+
Visibility, VisibilityMap)
88
from modelscope.hub.utils.aigc import AigcModel
99
from modelscope.hub.utils.utils import resolve_endpoint
1010
from modelscope.utils.constant import REPO_TYPE_MODEL, REPO_TYPE_SUPPORT
@@ -82,6 +82,20 @@ def define_args(parsers: _SubParsersAction):
8282
help=
8383
'If True, do not raise error when repo already exists. Defaults to False.',
8484
)
85+
parser.add_argument(
86+
'--gated',
87+
dest='gated_mode',
88+
action='store_true',
89+
default=None,
90+
help=
91+
'Enable gated mode (application-based download) for private repos.',
92+
)
93+
parser.add_argument(
94+
'--no-gated',
95+
dest='gated_mode',
96+
action='store_false',
97+
help='Disable gated mode for private repos (normal private).',
98+
)
8599
parser.add_argument(
86100
'--endpoint',
87101
type=str,
@@ -176,6 +190,7 @@ def _create_regular_repo(self):
176190
exist_ok=self.args.exist_ok,
177191
create_default_config=True,
178192
endpoint=endpoint,
193+
gated_mode=self.args.gated_mode,
179194
)
180195

181196
def _create_aigc_model(self):
@@ -225,7 +240,8 @@ def _create_aigc_model(self):
225240
visibility=visibility_idx,
226241
license=self.args.license,
227242
chinese_name=self.args.chinese_name,
228-
aigc_model=aigc_model)
243+
aigc_model=aigc_model,
244+
gated_mode=self.args.gated_mode)
229245
print(f'Successfully created AIGC model: {model_url}')
230246
except Exception as e:
231247
print(f'Error creating AIGC model: {e}')

‎modelscope/hub/api.py‎

Lines changed: 43 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -329,7 +329,8 @@ def create_model(self,
329329
original_model_id: Optional[str] = '',
330330
endpoint: Optional[str] = None,
331331
token: Optional[str] = None,
332-
aigc_model: Optional['AigcModel'] = None) -> str:
332+
aigc_model: Optional['AigcModel'] = None,
333+
gated_mode: Optional[bool] = None) -> str:
333334
"""Create model repo at ModelScope Hub.
334335
335336
Args:
@@ -343,6 +344,9 @@ def create_model(self,
343344
aigc_model (AigcModel, optional): AigcModel instance for AIGC model creation.
344345
If provided, will create an AIGC model with automatic file upload.
345346
Refer to modelscope.hub.utils.aigc.AigcModel for details.
347+
gated_mode (bool, optional): Gated mode for private repos.
348+
True = gated (application-based download), False = off (normal private).
349+
Only effective when visibility is PRIVATE (1).
346350
347351
Returns:
348352
str: URL of the created model repository
@@ -374,6 +378,12 @@ def create_model(self,
374378
'TrainId': os.environ.get('MODELSCOPE_TRAIN_ID', '')
375379
}
376380

381+
if gated_mode is not None:
382+
if visibility != ModelVisibility.PRIVATE:
383+
logger.warning('gated_mode is only effective when visibility is PRIVATE, ignored.')
384+
else:
385+
body['ProtectedMode'] = 1 if gated_mode else 2
386+
377387
# Set path based on model type
378388
if aigc_model is not None:
379389
# Use AIGC model endpoint
@@ -1428,7 +1438,8 @@ def create_dataset(self,
14281438
visibility: Optional[int] = DatasetVisibility.PUBLIC,
14291439
description: Optional[str] = '',
14301440
endpoint: Optional[str] = None,
1431-
token: Optional[str] = None) -> str:
1441+
token: Optional[str] = None,
1442+
gated_mode: Optional[bool] = None) -> str:
14321443
"""
14331444
Create a dataset in ModelScope.
14341445
@@ -1441,6 +1452,9 @@ def create_dataset(self,
14411452
description (str, optional): The description of the dataset. Defaults to ''.
14421453
endpoint (str, optional): The endpoint to use. If not provided, the default endpoint is used.
14431454
token (str, optional): The access token for authentication.
1455+
gated_mode (bool, optional): Gated mode for private repos.
1456+
True = gated (application-based download), False = off (normal private).
1457+
Only effective when visibility is PRIVATE (1).
14441458
14451459
Returns:
14461460
str: The URL of the created dataset repository.
@@ -1462,6 +1476,12 @@ def create_dataset(self,
14621476
'Description': (None, description)
14631477
}
14641478

1479+
if gated_mode is not None:
1480+
if visibility != DatasetVisibility.PRIVATE:
1481+
logger.warning('gated_mode is only effective when visibility is PRIVATE, ignored.')
1482+
else:
1483+
files['ProtectedMode'] = (None, 1 if gated_mode else 2)
1484+
14651485
r = self.session.post(
14661486
path,
14671487
files=files,
@@ -2200,6 +2220,7 @@ def create_repo(
22002220
exist_ok: Optional[bool] = False,
22012221
create_default_config: Optional[bool] = True,
22022222
aigc_model: Optional[AigcModel] = None,
2223+
gated_mode: Optional[bool] = None,
22032224
**kwargs,
22042225
) -> str:
22052226
"""
@@ -2217,6 +2238,9 @@ def create_repo(
22172238
In the format of `https://www.modelscope.cn` or 'https://www.modelscope.ai'
22182239
exist_ok (Optional[bool]): If the repo exists, whether to return the repo url directly.
22192240
create_default_config (Optional[bool]): If True, create a default configuration file in the model repo.
2241+
gated_mode (Optional[bool]): Gated mode for private repos.
2242+
True = gated (application-based download), False = off (normal private).
2243+
Only effective when visibility is ``private``.
22202244
**kwargs: The additional arguments.
22212245
22222246
Returns:
@@ -2256,6 +2280,7 @@ def create_repo(
22562280
aigc_model=aigc_model,
22572281
token=token,
22582282
endpoint=endpoint,
2283+
gated_mode=gated_mode,
22592284
)
22602285
if create_default_config:
22612286
with tempfile.TemporaryDirectory() as temp_cache_dir:
@@ -2290,6 +2315,7 @@ def create_repo(
22902315
visibility=visibility,
22912316
token=token,
22922317
endpoint=endpoint,
2318+
gated_mode=gated_mode,
22932319
)
22942320
print(f'New dataset created successfully at {repo_url}.', flush=True)
22952321

@@ -3882,7 +3908,8 @@ def set_repo_visibility(self,
38823908
repo_id: str,
38833909
repo_type: Literal['model', 'dataset'],
38843910
visibility: Literal['private', 'public'],
3885-
token: Union[str, None] = None
3911+
token: Union[str, None] = None,
3912+
gated_mode: Optional[bool] = None,
38863913
) -> dict:
38873914
"""
38883915
Set the visibility of a repo.
@@ -3893,6 +3920,9 @@ def set_repo_visibility(self,
38933920
visibility (Literal['private', 'public']): The visibility to set, `private` or `public`.
38943921
token (Union[str, None]): The access token. If None, will use the cookies from the local cache.
38953922
See `https://modelscope.cn/my/myaccesstoken` to get your token.
3923+
gated_mode (Optional[bool]): Gated mode for private repos.
3924+
True = gated (application-based download), False = off (normal private).
3925+
Only effective when visibility is ``private``.
38963926
38973927
Returns:
38983928
dict: The response from the server.
@@ -3907,6 +3937,10 @@ def set_repo_visibility(self,
39073937
visibility_code: int = visibility_map.get(visibility, 5)
39083938
cookies = self.get_cookies(access_token=token, cookies_required=True)
39093939

3940+
if gated_mode is not None and visibility != 'private':
3941+
logger.warning('gated_mode is only effective when visibility is private, ignored.')
3942+
gated_mode = None
3943+
39103944
if repo_type == REPO_TYPE_MODEL:
39113945
model_info = self.get_model(model_id=repo_id, token=token)
39123946
path = f'{self.endpoint}/api/v1/models/{repo_id}'
@@ -3916,11 +3950,15 @@ def set_repo_visibility(self,
39163950
first = tasks[0]
39173951
if isinstance(first, dict) and first:
39183952
model_tasks = first.get('name')
3953+
if gated_mode is not None:
3954+
pm = 1 if gated_mode else 2
3955+
else:
3956+
pm = model_info.get('ProtectedMode', 2)
39193957
payload = {
39203958
'ChineseName': model_info.get('ChineseName', ''),
39213959
'ModelFramework': model_info.get('ModelFramework', 'Pytorch'),
39223960
'Visibility': visibility_code,
3923-
'ProtectedMode': 2,
3961+
'ProtectedMode': pm,
39243962
'ApprovalMode': model_info.get('ApprovalMode', 2),
39253963
'Description': model_info.get('Description', ''),
39263964
'AigcType': model_info.get('AigcType', ''),
@@ -3947,7 +3985,7 @@ def set_repo_visibility(self,
39473985
path = f'{self.endpoint}/api/v1/datasets/{dataset_idx}'
39483986
payload = {
39493987
'Visibility': visibility_code,
3950-
'ProtectedMode': 2,
3988+
'ProtectedMode': (1 if gated_mode else 2) if gated_mode is not None else 2,
39513989
}
39523990
else:
39533991
raise ValueError(f'Invalid repo type: {repo_type}, supported repos: {REPO_TYPE_SUPPORT}')

‎modelscope/hub/constants.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,18 @@ class Visibility(object):
169169
PUBLIC = 'public'
170170

171171

172+
class GatedMode(object):
173+
"""Gated mode for private repositories.
174+
175+
Only effective when Visibility is PRIVATE.
176+
API payload key: ``ProtectedMode``.
177+
Values: True = gated (application-based download),
178+
False = off (normal private).
179+
"""
180+
GATED = True
181+
OFF = False
182+
183+
172184
VisibilityMap = {
173185
ModelVisibility.PRIVATE: Visibility.PRIVATE,
174186
ModelVisibility.INTERNAL: Visibility.INTERNAL,

0 commit comments

Comments
 (0)