File size: 1,209 Bytes
24e7fa4 | 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 | import unittest
import tempfile
from pathlib import Path
import torch
from safetensors.torch import save_file
from tools.maya_export.common import iter_ckpt_tensors
class TestCommon(unittest.TestCase):
def test_iter_ckpt_tensors_skips_prefix(self):
# tiny safetensors file in a temp dir with keys backbone.lm.model.layers.0.w and depth.x;
# skip backbone.lm.model.layers. -> only depth.x yielded, once
with tempfile.TemporaryDirectory() as td:
p = Path(td) / "model.safetensors"
tensors = {
"backbone.lm.model.layers.0.w": torch.tensor([1.0, 2.0]),
"depth.x": torch.tensor([3.0, 4.0, 5.0]),
"model.depth.y": torch.tensor([6.0]),
}
save_file(tensors, str(p))
res = list(iter_ckpt_tensors(skip_prefixes=("backbone.lm.model.layers.",), files=[str(p)]))
names = [name for name, _ in res]
self.assertIn("depth.x", names)
self.assertIn("depth.y", names) # model. prefix stripped
self.assertNotIn("backbone.lm.model.layers.0.w", names)
self.assertEqual(len(res), 2)
if __name__ == "__main__":
unittest.main()
|