Spaces:
Sleeping
Sleeping
File size: 2,335 Bytes
a0288c0 259de4c a0288c0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | from pathlib import Path
import tempfile
import unittest
from unittest import mock
class ArtifactTests(unittest.TestCase):
def test_checkpoint_path_rejects_absolute_parent_and_empty_paths(self):
import artifacts
for value in ("/tmp/model", "../model", "a/../../model", ""):
with self.subTest(value=value), self.assertRaises(ValueError):
artifacts.normalize_checkpoint_path(value)
def test_download_checkpoint_validates_required_layout(self):
import artifacts
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
checkpoint = root / "checkpoints/30000"
(checkpoint / "params").mkdir(parents=True)
(checkpoint / "assets/ur_demo").mkdir(parents=True)
(checkpoint / "assets/ur_demo/norm_stats.json").write_text("{}")
with mock.patch.object(artifacts, "snapshot_download", return_value=str(root)):
paths = artifacts.download_checkpoint("owner/model", "checkpoints/30000")
self.assertEqual(paths.checkpoint, checkpoint)
self.assertEqual(paths.norm_stats, checkpoint / "assets/ur_demo/norm_stats.json")
def test_download_checkpoint_accepts_named_asset_directory(self):
import artifacts
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
checkpoint = root / "3000"
(checkpoint / "params/ocdbt.process_0/d").mkdir(parents=True)
stats = checkpoint / "assets/F-Fer/ur-1/norm_stats.json"
stats.parent.mkdir(parents=True)
stats.write_text("{}")
with mock.patch.object(artifacts, "snapshot_download", return_value=str(root)):
paths = artifacts.download_checkpoint("owner/model", "3000")
self.assertEqual(paths.norm_stats, stats)
def test_download_checkpoint_reports_missing_weights(self):
import artifacts
with tempfile.TemporaryDirectory() as directory:
with mock.patch.object(artifacts, "snapshot_download", return_value=directory):
with self.assertRaisesRegex(FileNotFoundError, "params/|model.safetensors"):
artifacts.download_checkpoint("owner/model", "checkpoint")
if __name__ == "__main__":
unittest.main()
|