Skip to content
This repository was archived by the owner on Jun 4, 2025. It is now read-only.

Commit 78694d9

Browse files
authored
update Zoo stub loading for SparseZoo 1.1 refactor (#54)
1 parent 15ba9b7 commit 78694d9

File tree

1 file changed

+3
-3
lines changed

1 file changed

+3
-3
lines changed

src/transformers/hf_argparser.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
from pathlib import Path
2424
from typing import Any, Iterable, List, NewType, Optional, Tuple, Union
2525

26-
from sparsezoo import Zoo
26+
from sparsezoo import Model
2727

2828
from .utils.logging import get_logger
2929

@@ -263,8 +263,8 @@ def _download_dataclass_zoo_stub_files(data_class: DataClass):
263263

264264
logger.info(f"Downloading framework files for SparseZoo stub: {val}")
265265

266-
zoo_model = Zoo.load_model_from_stub(val)
267-
framework_file_paths = zoo_model.download_framework_files()
266+
zoo_model = Model(val)
267+
framework_file_paths = [file.path for file in zoo_model.training.default.files]
268268
assert framework_file_paths, "Unable to download any framework files for SparseZoo stub {val}"
269269
framework_file_names = [os.path.basename(path) for path in framework_file_paths]
270270
if "pytorch_model.bin" not in framework_file_names or ("config.json" not in framework_file_names):

0 commit comments

Comments
 (0)