68 lines
2.7 KiB
Python
68 lines
2.7 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"
|
|
TEMPLATE = (ROOT / "patches" / "chat_template" / "default_chat_template.jinja").read_text(encoding="utf-8")
|
|
WHITELIST = apply_patches.load_whitelist(ROOT / "patches" / "architectures" / "whitelist.txt")
|
|
ALIASES = apply_patches.load_aliases(ROOT / "patches" / "architectures" / "aliases.json")
|
|
|
|
|
|
class ChatTemplateTest(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 test_injects_missing_template(self) -> None:
|
|
result = apply_patches.apply_chat_template(self.model, TEMPLATE)
|
|
self.assertTrue(result["injected"])
|
|
data = json.loads((self.model / "tokenizer_config.json").read_text(encoding="utf-8"))
|
|
self.assertEqual(data["chat_template"], TEMPLATE)
|
|
|
|
def test_keeps_existing_template(self) -> None:
|
|
path = self.model / "tokenizer_config.json"
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
data["chat_template"] = "already-set"
|
|
path.write_text(json.dumps(data), encoding="utf-8")
|
|
result = apply_patches.apply_chat_template(self.model, TEMPLATE)
|
|
self.assertFalse(result["injected"])
|
|
kept = json.loads(path.read_text(encoding="utf-8"))
|
|
self.assertEqual(kept["chat_template"], "already-set")
|
|
|
|
def test_force_overwrites_template(self) -> None:
|
|
path = self.model / "tokenizer_config.json"
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
data["chat_template"] = "already-set"
|
|
path.write_text(json.dumps(data), encoding="utf-8")
|
|
result = apply_patches.apply_chat_template(self.model, TEMPLATE, force=True)
|
|
self.assertTrue(result["injected"])
|
|
|
|
def test_missing_tokenizer_config_fails(self) -> None:
|
|
(self.model / "tokenizer_config.json").unlink()
|
|
with self.assertRaises(apply_patches.PatchError):
|
|
apply_patches.apply_chat_template(self.model, TEMPLATE)
|
|
|
|
def test_preview_renders_roles(self) -> None:
|
|
text = apply_patches.preview_default_template(
|
|
[
|
|
{"role": "system", "content": "你是助手"},
|
|
{"role": "user", "content": "你好"},
|
|
]
|
|
)
|
|
self.assertIn("<|user|>\n你好<|end|>\n", text)
|
|
self.assertTrue(text.endswith("<|assistant|>\n"))
|
|
self.assertIn("message['role']", TEMPLATE)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|