66 lines
2.6 KiB
Python
66 lines
2.6 KiB
Python
import json
|
|
import shutil
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
import apply_patches
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
FIXTURE = ROOT / "tests" / "fixtures" / "fake-model"
|
|
WHITELIST = apply_patches.load_whitelist(ROOT / "patches" / "architectures" / "whitelist.txt")
|
|
|
|
|
|
class ArchitectureTest(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.model = Path(self.tmp.name) / "model"
|
|
shutil.copytree(FIXTURE, self.model)
|
|
|
|
def tearDown(self) -> None:
|
|
self.tmp.cleanup()
|
|
|
|
def write_architectures(self, names: list[str]) -> None:
|
|
path = self.model / "config.json"
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
data["architectures"] = names
|
|
path.write_text(json.dumps(data), encoding="utf-8")
|
|
|
|
def test_whitelist_passes(self) -> None:
|
|
result = apply_patches.apply_architectures(self.model, WHITELIST, {}, "fail")
|
|
self.assertEqual(result["architectures"], ["LlamaForCausalLM"])
|
|
self.assertFalse(result["changed"])
|
|
|
|
def test_unknown_architecture_fails(self) -> None:
|
|
self.write_architectures(["ExampleForCausalLM"])
|
|
with self.assertRaises(apply_patches.PatchError) as ctx:
|
|
apply_patches.apply_architectures(self.model, WHITELIST, {}, "fail")
|
|
self.assertIn("ExampleForCausalLM", str(ctx.exception))
|
|
|
|
def test_unknown_architecture_skip_keeps_name(self) -> None:
|
|
self.write_architectures(["ExampleForCausalLM"])
|
|
result = apply_patches.apply_architectures(self.model, WHITELIST, {}, "skip")
|
|
self.assertEqual(result["skipped"], ["ExampleForCausalLM"])
|
|
data = json.loads((self.model / "config.json").read_text(encoding="utf-8"))
|
|
self.assertEqual(data["architectures"], ["ExampleForCausalLM"])
|
|
|
|
def test_alias_rewrites_to_whitelist(self) -> None:
|
|
self.write_architectures(["LlamaForCausalLMAlias"])
|
|
result = apply_patches.apply_architectures(
|
|
self.model,
|
|
WHITELIST,
|
|
{"LlamaForCausalLMAlias": "LlamaForCausalLM"},
|
|
"fail",
|
|
)
|
|
self.assertEqual(result["rewritten"], ["LlamaForCausalLMAlias->LlamaForCausalLM"])
|
|
data = json.loads((self.model / "config.json").read_text(encoding="utf-8"))
|
|
self.assertEqual(data["architectures"], ["LlamaForCausalLM"])
|
|
|
|
def test_shipped_aliases_are_empty(self) -> None:
|
|
aliases = apply_patches.load_aliases(ROOT / "patches" / "architectures" / "aliases.json")
|
|
self.assertEqual(aliases, {})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|