Skip to content

Commit

Permalink
Expose repo security status field in ModelInfo (#2639)
Browse files Browse the repository at this point in the history
* expose repo security status in ModelInfo

* remove warning in test
  • Loading branch information
hanouticelina authored Oct 28, 2024
1 parent 1fbbb6e commit 01bdce3
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 11 deletions.
7 changes: 5 additions & 2 deletions src/huggingface_hub/hf_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -760,6 +760,8 @@ class ModelInfo:
List of spaces using the model.
safetensors (`SafeTensorsInfo`, *optional*):
Model's safetensors information.
security_repo_status (`Dict`, *optional*):
Model's security scan status.
"""

id: str
Expand Down Expand Up @@ -788,6 +790,7 @@ class ModelInfo:
siblings: Optional[List[RepoSibling]]
spaces: Optional[List[str]]
safetensors: Optional[SafeTensorsInfo]
security_repo_status: Optional[Dict]

def __init__(self, **kwargs):
self.id = kwargs.pop("id")
Expand Down Expand Up @@ -853,7 +856,7 @@ def __init__(self, **kwargs):
if safetensors
else None
)

self.security_repo_status = kwargs.pop("securityRepoStatus", None)
# backwards compatibility
self.lastModified = self.last_modified
self.cardData = self.card_data
Expand Down Expand Up @@ -2493,7 +2496,7 @@ def model_info(
Whether to set a timeout for the request to the Hub.
securityStatus (`bool`, *optional*):
Whether to retrieve the security status from the model
repository as well.
repository as well. The security status will be returned in the `security_repo_status` field.
files_metadata (`bool`, *optional*):
Whether or not to retrieve metadata for files in the repository
(size, LFS metadata, etc). Defaults to `False`.
Expand Down
13 changes: 4 additions & 9 deletions tests/test_hf_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -1832,21 +1832,16 @@ def test_model_info(self):
self.assertIsInstance(model, ModelInfo)
self.assertEqual(model.sha, DUMMY_MODEL_ID_REVISION_ONE_SPECIFIC_COMMIT)

# TODO; un-skip this test once it's fixed.
@unittest.skip(
"Security status is currently unreliable on the server endpoint, so this"
" test occasionally fails. Issue is tracked in"
" https://github.com/huggingface/huggingface_hub/issues/1002 and"
" https://github.com/huggingface/moon-landing/issues/3695. TODO: un-skip"
" this test once it's fixed."
)
def test_model_info_with_security(self):
# Note: this test might break in the future if `security_repo_status` object structure gets updated server-side
# (not yet fully stable)
model = self._api.model_info(
repo_id=DUMMY_MODEL_ID,
revision=DUMMY_MODEL_ID_REVISION_ONE_SPECIFIC_COMMIT,
securityStatus=True,
)
self.assertEqual(model.securityStatus, {"containsInfected": False})
self.assertIsNotNone(model.security_repo_status)
self.assertEqual(model.security_repo_status, {"scansDone": True, "filesWithIssues": []})

def test_model_info_with_file_metadata(self):
model = self._api.model_info(
Expand Down

0 comments on commit 01bdce3

Please sign in to comment.